LeetCode 378. Kth Smallest Element in a Sorted Matrix

LeetCode 378

先想到的是利用已经有的顺序，每次向右向下遍历，加入一个min堆，然后从min堆中取出一个数继续按照这个顺序遍历。

    def kthSmallest2(self, matrix: List[List[int]], k: int) -> int:
        N = len(matrix)
        h = []
        visited = set()
        collections._heapq.heappush(h, [matrix[0][0], 0, 0])
        visited.add((0,0))
        while k>0 and h:
            cval, posx, posy = collections._heapq.heappop(h)            
            k = k-1

            rightx, righty = posx+1, posy            
            if rightx < N and righty < N and (rightx, righty) not in visited:
                collections._heapq.heappush(h, [matrix[rightx][righty], rightx, righty])
                visited.add((rightx, righty))

            bottomx, bottomy = posx, posy +1
            if bottomx < N and bottomy < N and (bottomx, bottomy) not in visited:
                collections._heapq.heappush(h, [matrix[bottomx][bottomy], bottomx, bottomy])
                visited.add((bottomx, bottomy))
            
        return cval

复杂度接近于O(k*logk)，提交了之后发现其实速度并不快。由于有visited的控制。
如果不做这个控制，直接把所有的数都放入max堆中，把所有超过k个数的都pop掉，只留最小的k个数。

    def kthSmallest1(self, matrix: List[List[int]], k: int) -> int:
        N = len(matrix)
        h = []
        for i in range(0, N):
            for j in range(0, N):
                collections._heapq.heappush(h, -matrix[i][j])
                if len(h) > k:
                    collections._heapq.heappop(h)
        
        return - collections._heapq.heappop(h)

代码很简单，时间复杂度 O（NNlog（N*N-k）），在当前给的case情况下，跑的还是比较快的。
最后一个方法是二分查找，
在最小和最大中间取一个mid，然后查看小于这个mid的数有多少个，然后根据这个信息来看需要取左边还是右边。求小于mid的数的时候还可以继续使用二分查找，找到每行中小于这个mid的数。
参考（https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/627572/Share-my-Python-Binary-Search-method）

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:

        # binary search        
        def helper(thr,matrix,k):
            count = 0
            for i in range(len(matrix)):
                L,R = 0,len(matrix[i])-1
                while L<=R:
                    mid = (L+R)//2
                    if matrix[i][mid]<thr:
                        L = mid+1
                    else:
                        R = mid-1 
                count+=L
            return count<k                
        
        n = len(matrix)
        L = ans = matrix[0][0]
        R = matrix[n-1][n-1]
        
        while L<=R:
            mid = (L+R)//2
            if helper(mid,matrix,k):
                ans = mid
                L = mid+1
            else:
                R = mid-1
        return ans        

再补一个类似的题目，LeetCode 668

    def findKthNumber(self, m: int, n: int, k: int) -> int:
        left = 1
        right = m*n

        def count(val, m, n):
            c = 0
            for i in range(1, m+1):
                if i > val: break
                c += min(n, val // i)
            return c

        while left <= right:
            mid = left + ((right-left) >> 1)
            if count(mid, m, n) >= k:
                right = mid -1
            else:
                left = mid + 1

        return left