关于python numpy ndarray 遍历太慢的问题

Question

假设有三个矩阵，如下

a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[1,2,3],[4,5,6],[7,8,9]])
c = np.array([[1,0,6],[8,0,0],[0,2,5]])

想通过判断c矩阵相应位置的值，来决定a相应位置的值是否与b相应位置的值相加，直接想到的做法是遍历，如下：

for i in range(3):
    for j in range(3):
        if c[i][j]==0:
            a[i][j] = a[i][j] + b[i][j]

但当a,b,c都很大时运算相当慢，有没有这种运算的矢量写法来提升效率？不用for循环

Answer 1

使用where函数进行条件判断，代码这样写：

import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = np.array([[1,2,3],[4,5,6],[7,8,9]])
c = np.array([[1,0,6],[8,0,0],[0,2,5]])
d=np.where(c==0,a+b,a)
print(d)

运行结果：

[[ 1  4  3]
 [ 4 10 12]
 [14  8  9]]

如有帮助 ，请点采纳 。