python中二维上最快的迭代

Question

在 python 中迭代二维数组的最快方法是什么？考虑到我总是需要 x 和 y 索引。

例如，我有这段代码试图“匹配”矩阵中相同数字的 3 个，如果它们是数字 2、3、4、5 或 6：

for x in range(matrix.shape[0]):
    for y in range(matrix.shape[1]):
        if matrix[x][y] == 2 or matrix[x][y] == 3 or matrix[x][y] == 4 or matrix[x][y] == 5 or matrix[x][y] == 6: ## if i find one of the numbers i need
            fuseNumber = matrix[x][y] ## lets get that number
            if matrix[x+1][y] == fuseNumber: ## if we find another of that to the bottom of the initial one, we should try to find a third one
                if matrix[x-1][y] == fuseNumber: ## if we find another one at the top from the initial one
                    matrix[x][y] = 0
                    matrix[x+1][y] = 0
                    matrix[x-1][y] = 0
                    ...

代码将继续使用一些 if，就像这个一样，以确保它测试所有可能的组合，但这并不重要。 我尝试将其更改为：

it = numpy.nditer(matrix, flags=['multi_index'])
while not it.finished:
    x = it.multi_index[0]
    y = it.multi_index[1]
    if matrix[x][y] == 2 or matrix[x][y] == 3 or matrix[x][y] == 4 or matrix[x][y] == 5 or matrix[x][y] == 6:
        fuseNumber = matrix[x][y] ## lets get that fuse number, whichever it may be!
        if matrix[x+1][y] == fuseNumber: ## if we find another of that to the bottom of the initial one, we should try to find a third one
            if matrix[x-1][y] == fuseNumber: ## if we find another one at the top from the initial one  ///////  center bottom top
                matrix[x][y] = 0
                matrix[x+1][y] = 0
                matrix[x-1][y] = 0
                ...
it.iternext()

但是使用 timeit.timeit() 它向我展示了第二个代码实际上更慢。尽管有这两个示例，但您将如何编写相同的代码，但获得最佳性能？

谢谢！

Answer 1

而不是使用双 for 循环进行搜索：

for x in range(matrix.shape[0]):
    for y in range(matrix.shape[1]):
        if matrix[x][y] == 2 or matrix[x][y] == 3 or matrix[x][y] == 4 or matrix[x][y] == 5 or matrix[x][y] == 6: ## if i find one of the numbers i need
            fuseNumber = matrix[x][y] ## lets get that fuse number,
         ...

您可以使用 where 搜索 example

这允许仅在所需范围内的 x、y 值上循环。

for x, y in zip(*np.where(np.logical_and(matrix>=2, matrix<=6))):
    fuseNumber = matrix[x][y] ## lets get that fuse number,
         ...

测试

测试这两种方法是否创建相同的 x,y 对

# Random matrix (5, 10) with numbers between 0 & 10
matrix = np.random.randint(0, 10,  size=(5, 10))

# OP method
set_for = set()
for x in range(matrix.shape[0]):
    for y in range(matrix.shape[1]):
      if 2 <= matrix[x][y] <= 6:
        set_for.add((x, y))

# Method using where
set_numpy = set((x, y) for x, y in zip(*np.where(np.logical_and(matrix>=2, matrix<=6))))

# Result
print(set_for == set_numpy)

输出

True