python中2维上最快的迭代

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
                    ...

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

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,
         ...

您可以使用whereexample进行搜索

这只允许在所需范围内的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