k th smallest element in row wise and column wise sorted matrix

Find the kth smallest element in row wise and column wise sorted element:


Example:


matrix[row][col] < matrix[row+1][col]


matrix[row][col] < matrix[row][col+1]


mat=



1 2 21


11 23 25


31 36 41


I have written following code, How to make more readable and improve its complexity:

import heapq
def kth_smallest(k, mat):
    m = len(mat)
    n = len(mat[0])
    hp = []
    for i in xrange(n):
           heapq.heappush(hp,[mat[0][i], 0, i])
    heapq.heapify(hp)
    while k >1:
        x = heapq.heappop(hp)
        if x[1]+1 < m:
            heapq.heappush(hp,[mat[x[1]+1][x[2]], x[1]+1, x[2]])           
        heapq.heapify(hp)
        k-=1

    return hp[0][0]

print kth_smallest(8, [[1,10,12],[2,11,21], [3,31,45]])

• heappush and heappop maintain the heap invariant; there is no need to call heapify after them. Alternatively, you could initialize the heap from an unsorted list using heapify and skip the heappush, but as the list is already sorted in this case, it is already a heap and you don't need to apply either function.

• You push the whole first row to the heap but the first k elements would suffice thanks to the rows being already sorted.

-
Instead of

for i in xrange(n):
       heapq.heappush(hp,[mat[0][i], 0, i])

it would be more idiomatic to use

for i, value in enumerate(mat[0]):
       heapq.heappush(hp, [value, 0, i])

though as I already advised to take only k elements, you could instead use

for i, value in zip(xrange(k), mat[0]):

-
More idiomatic instead of while k >1: ... k-=1 would be for _ in xrange(k - 1):

-
Readability of
[mat[x[1]+1][x[2]], x[1]+1, x[2]]

could be improved to [mat[row + 1][col], row + 1, col]

by unpacking value, row, col = heapq.heappop(hp)

Revised code:

def kth_smallest(k, mat):
    '''Return the kth smallest element in row wise and column wise sorted matrix.
    mat represents the matrix as list of lists.
    '''
    m = len(mat)
    hp = [(value, 0, col) for col, value in zip(xrange(k), mat[0])]
    # hp is a min heap because mat[0] is assumed to be sorted
    for _ in xrange(k - 1):
        value, row, col = heapq.heappop(hp)
        row += 1
        if row < m:
            heapq.heappush(hp, (mat[row][col], row, col))           

    return hp[0][0]

StackExchange Code Review Q#161101, answer score: 4