有序矩阵中第K小的元素,与数组中的第K个最大元素算法思考

有序矩阵中第K小的元素,与数组中的第K个最大元素算法思考_第1张图片

使用大根堆,第K小的元素,保存所有元素于大根堆,大根堆的根节点为最大值;pop根节点,即将最大值删除,则剩下的k个元素中,首元素就是就是倒数第K大,即第K小;
数组中的第K个最大元素用的是小根堆,有序矩阵中第K小的元素用的是大根堆,思考一下这种算法?

这道题让我们求有序矩阵中第K小的元素,这道题的难点在于数组并不是蛇形有序的,意思是当前行的最后一个元素并不一定会小于下一行的首元素,所以我们并不能直接定位第K小的元素,所以只能另辟蹊径。先来看一种利用堆的方法,我们使用一个最大堆,然后遍历数组每一个元素,将其加入堆,根据最大堆的性质,大的元素会排到最前面,然后我们看当前堆中的元素个数是否大于k,大于的话就将首元素去掉,循环结束后我们返回堆中的首元素即为所求

import heapq
class Solution(object):
    def kthSmallest(self, matrix, k):
        row = len(matrix)
        if row<1:
            return
        col = len(matrix[0])
        
        heap = []
        for i in range(row):
            for j in range(col):#遍历矩阵中的每个元素,添加到大根堆
                heapq.heappush(heap, -matrix[i][j])#使用负数,表示大根堆
                if len(heap) > k:#如果>k,heapq.heappop(heap)将大根堆的首元素pop
                    heapq.heappop(heap)
        return -heap[0]

你可能感兴趣的:(数据结构-堆)