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.
- Create an empty min heap using heapq in python.
- Now assign first row (list) in result variable and convert result list into min heap using heapify method.
- Now traverse remaining row elements and push them into created min heap.
- Now get k’th smallest element using nsmallest(k, iterable) method of heapq module.
Implementation:
Python3
# Function to find K'th smallest element in# a 2D array in Pythonimport heapqdef 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 programif __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.
Python3
import heapqdef 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]+1mat = [[10, 25, 20, 40], [15, 45, 35, 30], [24, 29, 37, 48], [32, 33, 39, 50]]k = 7print("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.
