数据结构 | 堆 | 问题12

在堆中，查找第k小的元素是一个经典问题。给定一个无序的数组，要求找到第k小的元素。

解决方案

方法一：堆排序

堆排序的思想是先将数组建立为一个小根堆，然后取出堆顶元素，用一个计数器统计取出的元素个数，当计数器等于k时，即为所求的第k小元素。如果要找到前k小的元素，则可以将取出的元素加入一个结果数组中，当计数器等于k时，输出结果数组即可。

时间复杂度：O(nlogn)

def kthSmallest(nums, k):
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
    for i in range(k):
        res = heapq.heappop(heap)
    return res

方法二：快速选择

快速选择的思想是基于快速排序的划分思想来实现的，先把数组划分为两部分，一部分元素比划分元素小，一部分元素比划分元素大，然后根据划分元素和k的大小关系再对其中一部分进行划分。递归地进行这个过程，直到划分元素正好为第k小元素。

时间复杂度：平均时间复杂度O(n)，最坏情况下的时间复杂度是O(n^2)。

def kthSmallest(nums, k):
    def partition(left, right, pivot_index):
        pivot = nums[pivot_index]
        # 1. move pivot to end
        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]  
        store_index = left
        # 2. move all smaller elements to the left
        for i in range(left, right):
            if nums[i] < pivot:
                nums[store_index], nums[i] = nums[i], nums[store_index]
                store_index += 1
        # 3. 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 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)
        if k_smallest == pivot_index:
             return nums[k_smallest]
        elif k_smallest < pivot_index:
             return select(left, pivot_index - 1, k_smallest)
        else:
             return select(pivot_index + 1, right, k_smallest)

    return select(0, len(nums) - 1, k - 1)

总结

无论是堆排序还是快速选择，都需要进行元素比较和交换，时间复杂度主要取决于比较和交换的次数。在对全部元素进行排序的时候，堆排序的效率优于快速选择，因为堆排序中可以利用堆的性质减少比较和交换的次数，并且可以实现原地排序。但是在只需要找到第k小元素的场景中，快速选择是更好的选择。