Python solution O(klogk) similar to problem 373

  • 3
    class Solution(object):
        def kthSmallest(self, matrix, k):
            :type matrix: List[List[int]]
            :type k: int
            :rtype: int
            result, heap = None, []
            heapq.heappush(heap, (matrix[0][0], 0, 0))
            while k > 0:
                result, i, j = heapq.heappop(heap)
                if i == 0 and j + 1 < len(matrix):
                    heapq.heappush(heap, (matrix[i][j + 1], i, j + 1))
                if i + 1 < len(matrix):
                    heapq.heappush(heap, (matrix[i + 1][j], i + 1, j))
                k -= 1
            return result

    Since the matrix is sorted, we do not need to put all the elements in heap at one time. We can simply pop and put for k times. By observation, if we look at the matrix diagonally, we can tell that if we do not pop matrix[i][j], we do not need to put on matrix[i][j + 1] and matrix[i + 1][j] since they are bigger.

    e.g., given the matrix below:
    1 2 4
    3 5 7
    6 8 9
    We put 1 first, then pop 1 and put 2 and 3, then pop 2 and put 4 and 5, then pop 3 and put 6...

  • -1

    Nice. I did it in a similar way. Although, I believe the complexity should be O(nlogk) for an nxn matrix. At most, n elements in the heap.

  • 0

    @ichuen Yes, if n < k

  • 2

    @ichuen Sorry I am a beginner, but why it is O(nlogk) not O(klogk)?
    Since we must iterate k times with each time's (at most) complexity logk (from heappush)


  • 0

    as discussed in the most popular thread, when K is big, say k = O(n^2), the performance is bad.

    use the media-of-the-median algorithm, finding k-th element of an array is linear to the size of the elements. so the performance will be O(n^2)

Log in to reply

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