703. Kth Largest Element in a Stream

LeetCode 703. Kth Largest Element in a Stream

Description

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Implement KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums.

• int add(int val) Returns the element representing the kth largest element in the stream.

Example 1:

Input:
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output:
[null, 4, 5, 5, 8, 8]

Explanation
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

Constraints:

• 1 <= k <= 104

• 0 <= nums.length <= 104

• -104 <= nums[i] <= 104

• -104 <= val <= 104

• At most 104 calls will be made to add.

• It is guaranteed that there will be at least k elements in the array when you search for the kth element.

Tags

Heap, Design

Solution

Build a heap with capacity of k, where the root is the minimum element, to store the top-k largest elements in the given array.

• Constructor: we keep pushing elements into the heap until the size of the heap is k. Then, we compare the current element with the root of the heap, if the root is smaller, we pop an element from the heap and push the current one into it.

• Add: perform the same logic inside the for-loop of the constructor, and return the root of the heap.

Complexity

• Time complexity

  • Initialization: $O(n\log(k))$

  • Add: $O(\log(k))$

• Space complexity: $O(k)$

Code

type intHeap []int

func (h intHeap) Len() int {
	return len(h)
}

func (h intHeap) Less(i, j int) bool {
	return h[i] < h[j]
}

func (h intHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *intHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *intHeap) Pop() interface{} {
	x := (*h)[len(*h)-1]
	*h = (*h)[:len(*h)-1]
	return x
}

type KthLargest struct {
	k int
	intHeap
}

func Constructor(k int, nums []int) KthLargest {
	h := make(intHeap, 0, k)
	for _, num := range nums {
		if len(h) < k {
			heap.Push(&h, num)
		} else if num > h[0] {
			heap.Pop(&h)
			heap.Push(&h, num)
		}
	}
	return KthLargest{k, h}
}

func (this *KthLargest) Add(val int) int {
	if len(this.intHeap) < this.k {
		heap.Push(&this.intHeap, val)
	} else if val > this.intHeap[0] {
		heap.Pop(&this.intHeap)
		heap.Push(&this.intHeap, val)
	}
	return this.intHeap[0]
}