Python | 在矩阵中搜索元素

在Python中，可以使用各种数据结构来表示矩阵，例如列表（list）、NumPy数组、Pandas DataFrame等等。本文将以列表为例，介绍如何在Python中在矩阵中搜索元素。

列表中的矩阵

列表可以嵌套，因此可以用列表表示矩阵。例如下面的代码创建了一个3 x 3的矩阵：

matrix = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]

线性搜索

对于较小的矩阵，可以使用线性搜索的方式来查找元素。例如下面的代码可以在列表中搜索指定的元素：

def search(matrix, target):
    for row in matrix:
        if target in row:
            return True
    return False

print(search(matrix, 5))  # True
print(search(matrix, 10))  # False

线性搜索的时间复杂度是 $O(mn)$，其中 $m$ 和 $n$ 分别是矩阵的行数和列数。

二分搜索

如果矩阵是已排序的，也可以使用二分搜索的方式来查找元素。例如下面的代码可以在已排序的列表中搜索指定的元素：

def search(matrix, target):
    m = len(matrix)
    n = len(matrix[0])
    low = 0
    high = m * n - 1
    while low <= high:
        mid = (low + high) // 2
        row = mid // n
        col = mid % n
        if matrix[row][col] == target:
            return True
        elif matrix[row][col] < target:
            low = mid + 1
        else:
            high = mid - 1
    return False

matrix = [[1, 3, 5, 7],
          [10, 11, 16, 20],
          [23, 30, 34, 60]]
print(search(matrix, 3))  # True
print(search(matrix, 13))  # False

二分搜索的时间复杂度为 $O(\log(mn))$。

双指针搜索

对于部分有序的矩阵，也可以使用双指针的方式来查找元素。例如下面的代码可以在单调不降的矩阵中搜索指定的元素：

def search(matrix, target):
    m = len(matrix)
    n = len(matrix[0])
    row = 0
    col = n - 1
    while row < m and col >= 0:
        if matrix[row][col] == target:
            return True
        elif matrix[row][col] < target:
            row += 1
        else:
            col -= 1
    return False

matrix = [[1, 4, 7, 11, 15],
          [2, 5, 8, 12, 19],
          [3, 6, 9, 16, 22],
          [10, 13, 14, 17, 24],
          [18, 21, 23, 26, 30]]
print(search(matrix, 5))  # True
print(search(matrix, 20))  # False

双指针搜索的时间复杂度为 $O(m + n)$。

以上是在Python中在矩阵中搜索元素的几种常用方式，具体选择哪种方式取决于矩阵的大小、有序性等因素。