Shortest possible solution in Python (seriously!)


  • 5
    class Solution(object):
        def kthSmallest(self, matrix, k):
            return list(heapq.merge(*matrix)][k-1]
    

    Ok, kidding aside. This is O(klog(n)).

    class Solution(object):
        def kthSmallest(self, matrix, k):
            heap, res = [(row[0], r, 0) 
                        for r, row in enumerate(matrix) 
                        if row], 0
    
            heapq.heapify(heap)
            for k in range(1, k + 1):
                res, row, col = heapq.heappop(heap)
                if col + 1 < len(matrix[row]):
                    heapq.heappush(heap, (matrix[row][col + 1], row, col + 1))
            return res
    

    This is O(klog(k))

    class Solution(object):
        def kthSmallest(self, matrix, k):
            heap, res, n = [(matrix[0][0], 0, 0)], 0, len(matrix)
            for k in range(1, k + 1):
                res, row, col = heapq.heappop(heap)
                if not row and col < n - 1:
                    heapq.heappush(heap, (matrix[row][col + 1], row, col + 1))
                if row < n - 1:
                    heapq.heappush(heap, (matrix[row + 1][col], row + 1, col))
            return res
    

  • 0

    Since you're merging the complete matrix regardless of k, the complexity isn't O(klog(n)) but only O(n2log(n))l


  • 0

    @StefanPochmann heaps.merge returns an iterator, right? Getting one element from the iterator is O(log(n^2)) (so O(log(n)), I do k extractions, so it's O(klog(n)). What is the issue? You mean when k = n^2 I get n^2log(n)?


  • 0

    @agave said in Shortest possible solution in Python (seriously!):

    I do k extractions

    No, you do n2 extractions. Because list(...) takes everything. Only afterwards you apply k.

    Btw, time to fix your title, as @o_sharp posted a shorter Python solution than your "seriously! shortest possible" one :-)


  • 0

    @StefanPochmann oh yeah, you're right, I didn't pay attention! What about

    class Solution(object):
        def kthSmallest(self, matrix, k):
            x = heapq.merge(*matrix)
            return [next(x) for i in range(k)][k-1]
    

    And shush, don't tell people there is a shorter one around! :-)


  • 0

    Yes, that one is O(k log(n)). You could also use index -1 instead of k-1. Or itertools.islice instead of the list comprehension.


  • 0
    V

    @StefanPochmann
    x = heapq.merge(*matrix)
    return list( itertools.islice( x,k ) )[-1]


  • 0

    @VivianS Yeah, that's one way. Still takes O(n+k) space, though. Can you also do it with O(n) space?


  • 1

    @StefanPochmann what about this?

    class Solution(object):
        def kthSmallest(self, matrix, k):
            x = heapq.merge(*matrix)
            return list(itertools.islice(x, k-1, k))[-1]
    

  • 0

    @agave Yep, though I'd use next instead of creating+accessing a list.


  • 0
    V
    This post is deleted!

  • 0
    G

    Brute force is pretty short too:

    def kthSmallest(self, matrix, k):
        return sorted([ele for row in matrix for ele in row])[k-1]
    

    Runtime: 86ms


Log in to reply
 

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