leetcode/0789_kth-largest-element-in.../python3/naive_approach_tle.py

46 lines
1.2 KiB
Python
Raw Normal View History

# Naive approach, will TLE
# Time for each add: O(k)
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
len_nums = len(nums)
# Only keep max k elements, rest are useless
if len_nums >= k:
self.nums = sorted(nums)[len_nums - k:]
# If we're provided less than k elements, we
# will it up with least possible number
else:
self.nums = [
*([float("-inf")] * (k - len_nums)),
*sorted(nums)
]
def add(self, val: int) -> int:
# Find the right place for val in the list
# of length k
i = self.k - 1
while i >= 0 and val < self.nums[i]:
i -= 1
# If val is big enough to fit in the list, shift
# existing elements to left
j = 0
while j < i:
self.nums[j] = self.nums[j + 1]
j += 1
if i >= 0:
self.nums[i] = val
# kth largest element is always at index 0
return self.nums[0]
# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)