Kth Smallest Element in a Sorted Matrix / Heap Data Structure

Kth Smallest Element in a Sorted Matrix

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Review Questions:

1. What is heap?

2. What are some applications of heap data structure?

3. What are some useful functions in heapq module?

Heaps

Heap is a tree-based data structure in which the tree is a complete binary tree. Generally, heaps are either max-heap (node being the maximum value) or min-heap (1). Heapifying is a O(logN) operation while finding the max or min in max-heap/min-heap will be O(1).

Heap sort uses binary heap to sort an array in O(N logN) time, priority queues can be efficiently implemented using binary heap since it supports insert(), delete(), extractmax() all in O(logN) time, and these priority queues are used in graph algorithms such as Dijkstra’s Shortest Path and Prim’s Minimum Spanning Tree (2).

Useful Functions in ‘heapq’ module

• heapify(iterable) –> converts ‘iterable’ into a heap data structure
• heappush(heap, ele) –> inserts the element while keeping the order and heap data structure
• heappop(heap) –> removes and returns the smallest value from the heap while keeping the order and heap data structure
• heappushpop(heap, ele) –> combination of heappush and heappop
• heappreplace(heap,ele) –> similar to heappushpop, but heap pops an element first, meaning that it will return the smallest value in the original heap regardless of pushed element
• nlargest(k, iterable, key=fun) –> returns the k largest elements from the iterable specified and satisfying the key if mentioned

Approach 1:

Given that we want to maintain ‘N’ different variables with each of them pointing to an element in their corresponding list, heap data structure seems really useful.

1. We can first initialize a min-heap H. Since we need to find the kth smallest element, we only need to consider min(N,K) rows where N represents the number of rows.

2. We add min(N,K) elements and heapify the list.

3. At each step, we remove the minimum element from the heap, which tells us which row we need to add the next element(r,c+1), and iterate this process k times.

Implementation:

import heapq

class Solution(object):
def kthSmallest(self, matrix, k):

        N = len(matrix)
        minHeap = []
        for r in range(min(k,N)):
            minHeap.append((matrix[r][0], r, 0))

        heapq.heapify(minHeap)
        while k:
            element, r, c = heapq.heappop(minHeap)
            if c < N - 1:
                heapq.heappush(minHeap, (matrix[r][c+1], r, c+1))
            k -= 1
            
        return element

Complexity Analysis:

Time Complexity:O(X + K * logX) where X is min(N,k)

Space Complexity: O(X) which is occupied by the heap.

Approach 2:

For a sorted array, it would be really convenient to use Binary Search to find the Kth smallest number. However, this problem has a sorted matrix, not an array. If so, can we apply binary search on the number range instead of the index range by using a hypothetical middle element that will divide the matrix into left and right halves?

1. We start the binary search with start = matrix[0][0] and end = matrix[n-1][n-1].

2. Find the hypothtical middle element that can be, but not necessarily in our matrix.

3. Count all the numbers smaller than or equal to middle in the matrix. Since our matrix is sorted, this operation is O(N).

4. While counting, keep track of the smallest number greater than the middle- “R” and the biggest number less than or equal to the middle - “L”. We will use these numbers to adjust the number range for the binary search in the next iteration.

5. If the count is equal to K, L will be our required number.

6. If the count is less than K, we update start = R to search in the higher part of the matrix.

7. If the count is greater than K, we can update end = L to search in the lower part of the matrix.

Implementation:

class Solution(object):
    def countLessEqual(self, matrix, mid, smaller, larger):
        count, n = 0, len(matrix)
        row, col = n-1, 0
        
        while row >= 0 and col < n:
          if matrix[row][col] > mid:
              R = min(R, matrix[row][col])
              row -= 1
          else:
              L = max(L, matrix[row][col])
              count += row + 1
              col+= 1
        return count, L, R
    
    def kthSmallest(self,matrix,k):
        n = len(matrix)
        start, end = matrix[0][0], matrix[n-1][n-1]
        while start < end:
            mid = start + (end - start) / 2
            smaller, larger = (matrix[0][0], matrix[n - 1][n - 1])
            count, L, R = self.countLessEqual(matrix, mid, L, R)
            if count == k:
                return L
            if count < k:
                start = R
            else:
                end = L
        return start

Complexity Analysis:

Time Complexity: O(N * log(max - min))

Space Complexity: O(1) since binary search does not use any additional space.

Reflection:

This was a great exercise to review heap data structure as well as using Binary Search in an innovative way.

I feel like I need to be more familiar with the functions I can use for data structres provided by Python module!