LeetCode解題之Search a 2D Matrix

原題

在一個每行從左到右依次遞增，且下一行第一個數字比上一行最後一個數字大的矩陣中，判斷目標數字是否存在。

注意點：

例子：

輸入:

matrix = 
[
  [1,   3,  5,  7],
  [10, 11, 16, 20],
  [23, 30, 34, 50]
]
target = 3

輸出: True

解題思路

把矩陣從左到右、從上到下連起來就是一個遞增的數組，可以用二分搜索來查找。現在只要找出數組下標到矩陣的映射關系就可以了：i -> [i // n][i % n]，其中i是數組中的下標，n是矩陣的寬。

AC源碼

class Solution(object):
    def searchMatrix(self, matrix, target):
        """
        :type matrix: List[List[int]]
        :type target: int
        :rtype: bool
        """
        m = len(matrix)
        n = len(matrix[0])
        l, h = 0, m * n - 1
        while l <= h:
            mid = l + (h - l) // 2
            if matrix[mid // n][mid % n] == target:
                return True
            elif matrix[mid // n][mid % n] < target:
                l = mid + 1
            else:
                h = mid - 1
        return False


if __name__ == "__main__":
    assert Solution().searchMatrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]], 5) == True
    assert Solution().searchMatrix([[1, 2], [3, 4]], 4) == True
    assert Solution().searchMatrix([[1]], 2) == False