問題:
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.
heapではあるけど、１発目からそれにいくのはどうかな。とりあえずリストソートと考えちゃう。
1 リストをソート
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.top_k_scores = nums
self.top_k_scores.sort(reverse=True)
if k < len(self.top_k_scores):
self.top_k_scores = self.top_k_scores[:k]
def add(self, val: int) -> int:
self.top_k_scores.append(val)
self.top_k_scores.sort(reverse=True)
if self.k < len(self.top_k_scores):
self.top_k_scores = self.top_k_scores[:self.k]
return self.top_k_scores[-1]
時間計算量: O(nlogn) 空間計算量: O(n)
2 heapで実装
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.top_k_scores = [] #heap
for num in nums:
heapq.heappush(self.top_k_scores, num)
while self.k < len(self.top_k_scores):
heapq.heappop(self.top_k_scores)
def add(self, val: int) -> int:
heapq.heappush(self.top_k_scores, val)
while self.k < len(self.top_k_scores):
heapq.heappop(self.top_k_scores)
return self.top_k_scores[0]
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
heapq.heapify(nums)
self.top_k_scores = nums # heap
while self.k < len(self.top_k_scores):
heapq.heappop(self.top_k_scores)
def add(self, val: int) -> int:
heapq.heappush(self.top_k_scores, val)
while self.k < len(self.top_k_scores):
heapq.heappop(self.top_k_scores)
return self.top_k_scores[0]
heapの実装もしたほうがいいよ。できたらうpする。
余談。Quick Selectで無理やり実装
想定解ではないと思います。LeetCodeに投げたら 10/12 passedでTLEになりましたね。
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.top_k_scores = []
for num in nums:
self.top_k_scores.append(num)
def quick_select(self, left, right):
pivot_index = random.randint(left, right)
pivot = self.top_k_scores[pivot_index]
self.top_k_scores[right], self.top_k_scores[pivot_index] = self.top_k_scores[pivot_index], self.top_k_scores[right]
partition_index = left
for i in range(left, right):
if self.top_k_scores[i] > pivot:
self.top_k_scores[partition_index], self.top_k_scores[i] = self.top_k_scores[i], self.top_k_scores[partition_index]
partition_index += 1
self.top_k_scores[right], self.top_k_scores[partition_index] = self.top_k_scores[partition_index], self.top_k_scores[right]
if partition_index == self.k - 1:
return None
if partition_index < self.k - 1:
return self.quick_select(partition_index + 1, right)
if partition_index > self.k - 1:
return self.quick_select(left, partition_index - 1)
def add(self, val: int) -> int:
self.top_k_scores.append(val)
self.quick_select(0, len(self.top_k_scores) - 1)
return self.top_k_scores[self.k-1]