LeetCode 658. 找到 K 个最接近的元素 - Python
原题链接
中等
作者:
小聂同学
,
2023-10-09 01:22:35
,
所有人可见
,
阅读 47
class Solution:
# Find k elements cloesest to x
'''
Main Idea - Heap
Use a max heap to store tuple - (diffrence, val)
Time Complexity: O(nlogk), n is the number of list
Space Complexity: O(k)
'''
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
import heapq
res = []
max_heap = []
heapq.heapify(max_heap)
for num in arr:
heapq.heappush(max_heap, (-abs(num-x), -num))
if len(max_heap) > k:
heapq.heappop(max_heap)
while max_heap:
num = -1* heapq.heappop(max_heap)[1]
res.append(num)
res.sort()
return res
'''
Main idea - Binary Search + Two Pointer
Use binary search to find x position - O(logn)
Use Two pointers to find k elements meet the reqirement - O(k)
Time Complexity - O(logn +k)
Space Complexity - O(k)
'''
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
# Binary Search
left, right = 0, len(arr) -1
while left < right:
mid = (left + right +1) //2
if arr[mid] <= x:
left = mid
else:
right = mid -1
idx_x = left
# Corner Case:
# We find the first element <= x, but the first element > x may closer
if idx_x != len(arr)-1:
if abs(x-arr[idx_x]) > abs(x-arr[idx_x+1]):
idx_x = idx_x+1
# Two Pointer
left, right = idx_x, idx_x
while k-1 >0: # x already in the range
if left - 1 < 0:
right += 1
elif right +1 > len(arr)-1:
left -= 1
else:
if abs(x-arr[left-1]) > abs(x-arr[right+1]):
right +=1
else: # When equal select the small one
left -=1
k -=1
return arr[left:right+1]