该题解采用了快速选择(Quickselect)算法,这是一种用于在未排序的数组中找到第 k 个最小元素的方法。算法的核心思想是快速排序的一部分,即选取一个枢轴元素并对数组进行分区,将小于枢轴的元素放在其左侧,大于枢轴的元素放在其右侧。根据枢轴的位置与 k 的比较,递归地在需要的一侧继续进行分区操作,直到找到第 k 个最小元素。在本题中,我们不需要完全排序,只需要找到最小的 k 个元素即可。
时间复杂度: O(n)
空间复杂度: O(log n)
class Solution:
def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
if k == 0:
return [] # 如果 k 为 0,直接返回空列表
if len(arr) <= k:
return arr # 如果数组长度不大于 k,直接返回整个数组
self.partition_helper(arr, 0, len(arr) - 1, k) # 调用辅助函数来进行分区
return arr[:k] # 返回最小的 k 个数
def partition_helper(self, arr, lo, hi, k):
j = self.partition(arr, lo, hi) # 执行一次分区操作
if j > k:
self.partition_helper(arr, lo, j-1, k) # 如果枢轴位置大于 k,递归左侧
elif j < k:
self.partition_helper(arr, j+1, hi, k) # 如果枢轴位置小于 k,递归右侧
def partition(self, arr, lo, hi):
if lo == hi:
return lo
v = arr[lo] # 选择枢轴元素
i = lo
j = hi+1
while True:
while True:
i += 1
if i >= hi or arr[i] > v:
break
while True:
j -= 1
if j <= lo or arr[j] < v:
break
if i >= j:
break # 当两个指针相遇时跳出循环
arr[i], arr[j] = arr[j], arr[i] # 交换元素
arr[lo], arr[j] = arr[j], arr[lo] # 将枢轴元素放到正确位置
return j # 返回枢轴的最终位置