378. Kth Smallest Element in a Sorted Matrix (Python)

Related Topic

Heap.

Description

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Sample I/O

Example 1

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note

You may assume k is always valid, 1 ≤ k ≤ n2.

Methodology

Use Heap to solve the question.

Min Heap (easy understand): Push all values into min heap and pop k-1 value from min heap, then the head of the heap will be the kth element.

Max Heap (better perfermance): Push values into max heap and keep the max heap length equal to k, then the head of the heap will be the kth element.

Min Heap Solution:

def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        minHeap = []
        for x in matrix:
            for y in x:
                heapq.heappush(minHeap, y)
        index = 0
        while index<k-1:
            heapq.heappop(minHeap)
            index+=1
        return minHeap[0]

BigO

Time complexity : O($n^2$logn + klogk) where n is the size of input and k is the size of the kth largest element. Sum up to O($n^2$logn)

Max Heap Solution:

def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        maxHeap = []
        for x in matrix:
            for y in x:
                if len(maxHeap)<k:
                    heapq.heappush(maxHeap, -y)
                else:
                    heapq.heappushpop(maxHeap, -y)
        return -maxHeap[0]

BigO

Time complexity : O($n^2$logn) where n is the size of input.