Numpy 循环前进的方式

Question

我想在下面的代码中避免for-loop with step 并用Numpy代码替换它以加快进程：

import numpy as np

A = np.array([[0,0,7,0,0,0], [0,0,0,0,5,0]])

for i in range(A.shape[0]):
    for j in range(A.shape[1]-1):
        if A[i,j]==7 and A[i,j+1]==0:
            A[i,j+1]=7

我知道如何使用 for-loop 无步骤地做到这一点。比如说，A,B,C 是具有相同大小的二维数组，那么这个慢代码：

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        if A[i,j]==7 and B[i,j]==0:
            C[i,j]=7

...可以通过以下单行 numpy 代码更快：

C[(A==7) & (B==0)]=7

我猜应该是类似的东西，包括np.where和np.roll函数？感谢您的帮助！

Answer 1

我很确定有更好的方法，但如果您没有找到，这里有一个更快的方法：

for i in range(A.shape[1]-1):
  A = np.where((A==0)&(np.pad(A,((0,0),(1,0)))[:,:-1]==7),7,A)

输出：

[[0 0 7 7 7 7]
 [0 0 0 0 5 0]]

Numpy 循环前进的方式

Numpy way of step forward for loop

python

arrays

for-loop

numpy

broadcast