Share My Python Solution using Heap

  • 6


    1. h: list of tuple (element, row, index), which is initialised with first element of each row in the matrix.
    2. We maintain a heap. In the for loop, we get the smallest element v which is in row r, and replace v with the next element in the row r

    Time Complexity

    • insert an element into heap: O(log(n)), where n is the width of the matrix
    • find k the k-th element O(k)
    • Overall: O(klog(n))
    from heapq import heappush, heappop, heapreplace, heapify
    class Solution(object):
        def kthSmallest(self, matrix, k):
            :type matrix: List[List[int]]
            :type k: int
            :rtype: int
            h = [(row[0], row, 1) for row in matrix]
            # Since we want to find kth, we pop the first k elements 
            for _ in xrange(k - 1):
                v, r, i = h[0]
                if i < len(r):
                    heapreplace(h, (r[i], r, i + 1))
            return h[0][0]

  • 0

    This code is cool but it leads to log(n) heap operations which is supposed to be O(k).

  • 0

    This looks like O(N + KlogN) to me...

    I think you're forgetting the first step in your algorithm which is O(N) to initialize a heap in linear time. While not important if N and K are close, the big O(N) part can become dominate for very large N and very small K.

  • 0

    how can you heapq.heapify a list of tuples? does the key become the first element of each tuple?

  • 0

    and another question: isn't heapify(h) redundant?
    since the first column is always going to be a minheap anyway.

  • 0

    Sorry to write careless code. Like @StefanPochmann pointed below, you should have something like this instead:

    end_row = min(len(matrix), k)
    h = [(matrix[i][0], i, 1) for i in range(end_row)]

    The rest of the code needs some minor changes too, but the logic remains the same. The time complexity will be O(k log(min(k, n)))


    You are right for the code above, but we can make small modification as below to ensure the overall time complexity to be O(m log(m)) where m = min(k, n)

    end_row = min(len(matrix), k)
    h = [(row[0], row, 1) for row in matrix[:end_row]]

  • 1

    Actually the complexity is worse, because so far the comparison of heap elements has been ignored. If the matrix has the same number everywhere, then the element and the row in the (element, row, index) triples are always the same and their comparison causes another factor n. So O(kn log n). Demo:

    from timeit import timeit
    prev = None
    for e in range(7, 12):
        n = 2**e
        matrix = [[0] * n for _ in range(n)]
        k = n*n
        t = timeit(lambda: Solution().kthSmallest(matrix, k), number=1)
        print '%7.3f seconds  =>  factor %s' % (t, prev and t / prev)
        prev = t

    That doubles n from one iteration to the next, and you can see the time roughly increases by factor 8, not 4:

      0.027 seconds  =>  factor None
      0.170 seconds  =>  factor 6.20691851528
      1.285 seconds  =>  factor 7.56868340876
     10.324 seconds  =>  factor 8.03231050283
     88.901 seconds  =>  factor 8.61075319947

    But it can easily be turned into O(k log n) by storing (element, rownumber, row, index) tuples.

  • 0

    @StefanPochmann I updated my reply above. I think we could achieve O(klog(m)) where m = min(n, k), right?

    Initially O(m) to generate heap
    Then k of heap operations: O(klog(m))
    Overval: O(klog(m))

Log in to reply

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.