LeetCode 703. Kth Largest Element in a Stream (Python3)
はじめに
LeetCode 「703. Kth Largest Element in a Stream」の問題を Python3 で解きました。
問題
問題文
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.
和訳
あなたは大学の入学事務局の一員であり、志願者から提出されるテストスコアのうち、リアルタイムでk番目に高いスコアを追跡する必要があります。これにより、新たな志願者がスコアを提出するたびに、面接や入学の合格基準点を動的に決定することが可能になります。
与えられた整数 k に対して、試験スコアのストリームを維持し、新しいスコアが提出された後に k 番目に高いスコアを継続的に返すクラスを実装する任務が与えられています。具体的には、全スコアのソート済みリストにおける k 番目に高いスコアを求めています。
KthLargestクラスの実装:
KthLargest(int k, int[] nums)
整数kとテストスコアのストリームnumsでオブジェクトを初期化します。
int add(int val)
新しいテストスコアvalをストリームに追加し、これまでのテストスコアプールにおけるk番目に大きい要素を表す要素を返します。
例
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
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
制約
- 0 <= nums.length <= 10⁴
- 1 <= k <= nums.length + 1
- -10⁴ <= nums[i] <= 10⁴
- -10⁴ <= val <= 10⁴
- At most 10⁴ calls will be made to add.
解答1（bisect.insort）
add関数でリストにデータを挿入し、その上でk番目に大きい値を返します。
まずはわかりやすい方法として、ソート済みのリストを保持して、そこに要素を挿入します。
解答2よりも効率は悪いです。
bisect.insort()を使用します。
bisectライブラリは、ソート済みのリストを保ったまま新しい要素を挿入するための機能を提供しています。
挿入位置を二分探索で見つけます。
コード1
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.nums = sorted(nums)
def add(self, val: int) -> int:
bisect.insort(self.nums, val)
return self.nums[-(self.k)]
計算量
- 時間的計算量：要素追加ごとO(n)
- 挿入にかかるシフト＋保持する要素数分
- 空間的計算量：O(n)
- 全ての要素を保持
コード2
上記のコードは要素を全て保持していましたが、k番目に大きい要素を求めたいので、k個の大きい要素だけをソートして保持する方法を取ることもできます。
そうすると、リストの最小値がちょうどk番目に大きい要素となります。
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.nums = sorted(nums)[-k:]
def add(self, val: int) -> int:
if len(self.nums) < self.k:
bisect.insort(self.nums, val)
elif val > self.nums[0]:
bisect.insort(self.nums, val)
self.nums.pop(0)
return self.nums[0]
計算量
- 時間的計算量：要素追加ごとO(k)
- 挿入にかかるシフト＋保持する要素数分
- 空間的計算量：O(k)
- k個の要素だけ保持
解答2（heap）
次はheapを使ったより効率の良い方法です。
heap(ヒープ)は、データ構造の一種で、木構造のうち親要素が子要素より常に大きい(あるいは小さい)という条件を満たすものです。
Python の heapq は 最小値をすぐに取り出せるデータ構造（最小ヒープ）です。
k番目に大きい要素を求めたいので、常に大きい要素だけヒープに残すようにします。
そうすると、ヒープの先頭（最小値が）k番目に大きい要素になります。
リストをその都度ソートしたり挿入位置をずらすよりも、ヒープなら効率よく更新できます。
コード
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.heap = nums
heapq.heapify(self.heap)
while len(self.heap) > k:
heapq.heappop(self.heap)
def add(self, val: int) -> int:
if len(self.heap) < self.k:
heapq.heappush(self.heap, val)
elif val > self.heap[0]:
heapq.heapreplace(self.heap, val)
return self.heap[0]
初期化では、ヒープを作成し、k番目以降の要素だけを残しています。
add関数では、要素数がkより少なければ追加、valがヒープの最小値よりも大きければheapq.heapreplaceで入れ替えを行なっています。
計算量
- 時間的計算量：要素追加ごとO(log k)
- 空間的計算量：O(k)
- k個の要素だけ保持
Discussion