動機

原來binary search可以這麼用

Problem

Given an n x n matrix where each of the rows and columns are sorted in ascending order, return the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

 

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8Output: 13Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1Output: -5

 

Constraints:

  • n == matrix.length
  • n == matrix[i].length
  • 1 <= n <= 300
  • -109 <= matrix[i][j] <= 109
  • All the rows and columns of matrix are guaranteed to be sorted in non-decreasing order.
  • 1 <= k <= n2

Sol1: divide-and-conquer

當初想說對半折就好，但是這會把應該找的也一起砍掉

def g(val):
    return (val+(0 if val % 2 == 0 else 1))//2
def f(ms,k,i,j,li,lj):
    #print(k,i,j,li,lj)
    if k <= 1:
        return ms[i][j]
    else:
        if k >= (li*lj)//2:
            if li > 1 and lj > 1:
                if ms[i][j+1] <= ms[i+1][j]:
                    tmp = g(li)
                    k -= min(1,li-tmp)*lj
                    return f(ms, k ,i+li//2 ,j , tmp , lj)
                else:
                    tmp = g(lj)
                    k -= (lj-tmp)*li
                    return f(ms, k ,i ,j+lj//2 , li, tmp)
            elif li > 1:
                tmp = g(li)
                k -= (li-tmp)*lj
                return f(ms, k ,i+li//2 ,j , tmp , lj)
            else:
                tmp = g(lj)
                k -= (lj-tmp)*li
                return f(ms, k ,i ,j+lj//2 , li, tmp)
        else:
            if li > 1 and lj > 1:
                if ms[i][j+1] <= ms[i+1][j]:
                    tmp = g(li)
                    return f(ms, k ,i , j, tmp , lj)
                else:
                    tmp = g(lj)
                    return f(ms, k ,i , j, li, tmp)
            elif li > 1:
                tmp = g(li)
                return f(ms, k ,i , j, tmp , lj)
            else:
                tmp = g(lj)
                return f(ms, k ,i , j, li, tmp)
class Solution:
    def kthSmallest(self, ms: List[List[int]], k: int) -> int:
        return f(ms,k,0,0,len(ms),len(ms))

Sol2: binary search + binary search

第一個binary search在matrix的最大與最小的範圍中找個數字，依據數字的排名做二分 第二個binary search在每個row找總共有多少數字小於等於(upper bound)現在猜的數字

def count(ms,v):
    ret = 0
    exist = False
    for m in ms:
        tmp = bisect.bisect(m,v)
        ret += tmp
        if tmp > 0:
            exist = exist or m[tmp-1] == v
    return [ret, exist]
def check(k,n,exist):
    return n>=k and exist
class Solution:
    def kthSmallest(self, ms: List[List[int]], k: int) -> int:
        i = ms[0][0]
        j = ms[-1][-1]
        exist = False
        while i < j:
            mid = i + (j-i)//2
            tmp, exist = count(ms,mid)
            if tmp >= k:
                j = mid
            else:
                i = mid+1
        return i

Sol3: binary search + saddleback

利用saddleback的想法，從左下走

如果比猜的數字大，往上以及加總這排的小於等於的總數 如果比猜的數字小，往右走

不過要注意，不要往右走到超過邊界

def count(ms,v):
    ret = 0
    i = len(ms)-1
    j = 0
    while 0 <= i:
        if j == len(ms) or ms[i][j] > v:
            ret += j
            i -= 1
        else:
            j += 1 
    return ret
def check(k,n,exist):
    return n>=k and exist
class Solution:
    def kthSmallest(self, ms: List[List[int]], k: int) -> int:
        i = ms[0][0]
        j = ms[-1][-1]
        exist = False
        while i < j:
            mid = i + (j-i)//2
            tmp = count(ms,mid)
            if tmp >= k:
                j = mid
            else:
                i = mid+1
        return i