Python heapq to find K'th smallest element in a 2D array

Given an n x n matrix and integer k. Find the k'th smallest element in the given 2D array. Examples:

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

We will use similar approach like K’th Smallest/Largest Element in Unsorted Array to solve this problem.

1. Create an empty min heap using heapq in python.
2. Now assign first row (list) in result variable and convert result list into min heap using heapify method.
3. Now traverse remaining row elements and push them into created min heap.
4. Now get k'th smallest element using nsmallest(k, iterable) method of heapq module.

Implementation:

# 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)

7 th smallest element is  30

Approach#2: Using a Max Heap

We can create a max heap and keep adding the first row of elements in it. As we need to find the kth smallest element, we pop k-1 elements from the heap and add the next element from the same column as the popped element. The kth smallest element will be at the top of the heap after this process.

Algorithm

1. Create an empty max heap.
2. Push first row of elements in the heap.
3. Pop k-1 smallest elements from the heap and add next element from the same column as the popped element.
4. Return the kth smallest element.

import heapq

def kthSmallest(mat, k):
    # create a max heap with first row of elements
    heap = []
    for i in range(len(mat[0])):
        heapq.heappush(heap, (-mat[0][i], 0, i))
    # pop k-1 smallest elements from the heap
    for i in range(k-1):
        num, row, col = heapq.heappop(heap)
        # add next element from the same row
        if row < len(mat)-1:
            heapq.heappush(heap, (-mat[row+1][col], row+1, col))
    # the kth smallest element is now at the top of the heap
    return -heapq.heappop(heap)[0]+1

mat = [[10, 25, 20, 40],
       [15, 45, 35, 30],
       [24, 29, 37, 48],
       [32, 33, 39, 50]]
k = 7

print("Kth smallest element is:", kthSmallest(mat, k))

Kth smallest element is: 30

Time complexity: O(klog(n)), where n is the number of rows in the matrix. We need to pop k-1 elements and add at most n-1 elements for each popped element in the worst case. So, the time complexity is proportional to klog(n).

Space complexity: O(n), where n is the number of columns in the matrix. We are keeping a max heap with elements from the first row of the matrix. The maximum number of elements we need to store is equal to the number of columns in the matrix.