Kth Smallest Element in a Sorted Matrix https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/
- Build a min heap which takes O(n) time. Heapify k times which takes O(kLogn) time. Therefore, overall time complexity is O(n + kLogn) time. The above code can be optimized to build a heap of size k when k is smaller than n. In that case, the kth smallest element must be in first k rows and k columns.
- You can start with matrix[0,0] and at each iteration add the element to right and the element down. Notice that you might end up with duplicate elements and you will need a set to dedupe using a set.
- You can use this insight to avoid using a set:
Whenever an element is poll() from heap, push the element below it to heap, only push the element right to it to heap when it's in first row. So we can avoid duplicates.
from heapq import heappush, heappop class Solution(object): def kthSmallest(self, matrix, k): """ :type matrix: List[List[int]] :type k: int :rtype: int """ N = len(matrix) heap =  heappush(heap, (matrix, (0,0))) min_cnt = 1 while True: x, y = heappop(heap) if min_cnt == k: return x else: if y == 0: c = [(y+1, y), (y, y+1)] else: c = [(y+1, y)] for c1, c2 in c: if c1<N and c2<N: heappush(heap, (matrix[c1][c2], (c1,c2))) min_cnt = min_cnt + 1