LeetCode: Kth Largest Element in a Stream

You are part of a university admissions office and need to keep track of the kth highest test score from applicants in real-time. This helps to determine cut-off marks for interviews and admissions dynamically as new applicants submit their scores.

You are tasked to implement a class which, for a given integer k, maintains a stream of test scores and continuously returns the kth highest test score after a new score has been submitted. More specifically, we are looking for the kth highest score in the sorted list of all scores.

Implement the KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of test scores nums.
• int add(int val) Adds a new test score val to the stream and returns the element representing the kth largest element in the pool of test scores so far.

 

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

Example 2:

Input:
["KthLargest", "add", "add", "add", "add"]
[[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]

Output: [null, 7, 7, 7, 8]

Explanation:

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

 

Constraints:

• 0 <= nums.length <= 10^4
• 1 <= k <= nums.length + 1
• -10^4 <= nums[i] <= 10^4
• -10^4 <= val <= 10^4
• At most 10^4 calls will be made to add

 

 

 

이 문제는 주어진 list 중 k번째로 큰 수를 반환하는 문제이다. 

 

정렬 알고리즘을 사용하면 숫자를 추가할 때마다 다시 정렬을 해야한다는 단점이 있다. 

 

따라서 heap을 사용하는 것이 좀 더 적절한 문제이다. 

 

문제를 푸는 방식은 생각보다 간단하다. 하지만 생각하기 전에는 간단하지 않은 것 같다..

 

일단 k번쨰로 큰 수를 구하기 위해서는 최소 힙을 사용할 수 있다. 

 

주어진 nums 리스트를 heapify를 통해 최소 heap으로 만든 뒤에 heappop을 통해 원소 k개만 남기면 된다. 

 

가장 작은 수가 계속 삭제되기 때문에 k개 남았을때 heap의 root는 k번째로 큰 수가 되어서 

root를 반환하면 된다.

 

 

---

구현코드

class KthLargest(object):

    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        self.len = k
        self.heap = nums
        heapq.heapify(self.heap)

        while len(self.heap) > k:
            heapq.heappop(self.heap)
        

    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        
        heapq.heappush(self.heap, val)

        if len(self.heap) > self.len:
            heapq.heappop(self.heap)

        return self.heap[0]


# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)