# Share My Python Solution using Heap

• Algorithm

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]
heapify(h)

# 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))
else:
heappop(h)

return h[0][0]
``````

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

• 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.

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

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

• Updated:
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)))

Original:
@amaximov

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]]
heapify(h)
``````

• 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.

• @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))

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