题目
在未排序的数组中找到第 k 个最大的元素。请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。
code
- 随机选择一个枢轴。
- 使用划分算法将枢轴放在数组中的合适位置 pos。将小于枢轴的元素移到左边,大于等于枢轴的元素移到右边。
- 比较 pos 和 N - k 以决定在哪边继续递归处理。
在这里,由于知道要找的第 N - k
小的元素在哪部分中,我们不需要对两部分都做处理,这样就将平均时间复杂度下降到 O(N)。
import random
class Solution:
def findKthLargest(self, nums, k):
def partition(left, right, pivot_index):
pivot = nums[pivot_index]
# move pivot to end
nums[pivot_index], nums[right] = nums[right], nums[pivot_index]
# move all smaller elements to the left
store_index = left
for i in range(left, right):
if nums[i] < pivot:
nums[store_index], nums[i] = nums[i], nums[store_index]
store_index += 1
# move pivot to its final place
nums[right], nums[store_index] = nums[store_index], nums[right]
return store_index
def select(left, right, k_smallest):
# if the list contains only one element,return that element
if left == right:
return nums[left]
# select a random pivot_index between
pivot_index = random.randint(left, right)
# find the pivot position in a sorted list
pivot_index = partition(left, right, pivot_index)
# the pivot is in its final sorted position
if k_smallest == pivot_index:
return nums[k_smallest]
# go left
elif k_smallest < pivot_index:
return select(left, pivot_index - 1, k_smallest)
# go right
else:
return select(pivot_index + 1, right, k_smallest)
# kth largest is (n-k)th smallest
return select(0, len(nums) - 1, len(nums) - k)
a = Solution()
kth_largest = a.findKthLargest(nums=[3,2,1,5,6,4],k=2)
print(kth_largest)