703. Kth Largest Element in a St

Title: 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.

Your KthLargest class will have a constructor which accepts an integer k and an integer array nums, which contains initial elements from the stream. For each call to the method KthLargest.add, return the element representing the kth largest element in the stream.

Example:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3); // returns 4
kthLargest.add(5); // returns 5
kthLargest.add(10); // returns 5
kthLargest.add(9); // returns 8
kthLargest.add(4); // returns 8

Difficulty:

Easy

Implement Programming Language:

Go

Answer:

这里用的是最大堆实现，思路就是维护一个最大堆。Go实现堆需要实现一组接口就可以了。

代码：

package main

import "container/heap"

// copied from golang doc
// mininum setup of integer min heap
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{}) {
    // Push and Pop use pointer receivers because they modify the slice's length,
    // not just its contents.
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// real thing starts here
type KthLargest struct {
    Nums *IntHeap
    K    int
}

func Constructor(k int, nums []int) KthLargest {
    h := &IntHeap{}
    heap.Init(h)
    // push all elements to the heap
    for i := 0; i < len(nums); i++ {
        heap.Push(h, nums[i])
    }
    // remove the smaller elements from the heap such that only the k largest elements are in the heap
    for len(*h) > k {
        heap.Pop(h)
    }
    return KthLargest{h, k}
}

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