Python heapq 在二维数组中查找第 K 个最小元素

给定一个 nxn 矩阵和整数 k。在给定的二维数组中找到第 k 个最小的元素。

例子：

Input : mat = [[10, 25, 20, 40],
               [15, 45, 35, 30],
               [24, 29, 37, 48],
               [32, 33, 39, 50]]
        k =  7 
Output : 7th smallest element is 30

我们将使用类似未排序数组中第 K 个最小/最大元素的方法来解决这个问题。

1. 在Python使用 heapq 创建一个空的最小堆。
2. 现在分配结果变量中的第一行（列表）并使用heapify方法将结果列表转换为最小堆。
3. 现在遍历剩余的行元素并将它们推送到创建的最小堆中。
4. 现在使用 heapq 模块的nsmallest(k, iterable)方法获取第 k 个最小元素。

# Function to find K'th smallest element in 
# a 2D array in Python
import heapq
  
def kthSmallest(input):
  
    # assign first row to result variable
    # and convert it into min heap
    result = input[0] 
    heapq.heapify(result)
  
    # now traverse remaining rows and push
    # elements in min heap
    for row in input[1:]:
         for ele in row:
              heapq.heappush(result,ele)
  
    # get list of first k smallest element because
    # nsmallest(k,list) method returns first k 
    # smallest element now print last element of 
    # that list
    kSmallest = heapq.nsmallest(k,result)
    print (k,"th smallest element is ",kSmallest[-1])
  
# Driver program
if __name__ == "__main__":
    input = [[10, 25, 20, 40],
           [15, 45, 35, 30],
           [24, 29, 37, 48],
           [32, 33, 39, 50]]
    k = 7
    kthSmallest(input)

输出：

7th smallest element is 30

如果您希望与专家一起参加现场课程，请参阅DSA 现场工作专业课程和学生竞争性编程现场课程。