python - numpy.nditer 的正确用法？

Question

我正在尝试用 numpy.nditer 进行数组运算, 但没有得到预期的结果。

我的代码是

import numpy as np
arr1 = - np.random.random((2,2))
arr2 = np.random.random((2,2))
arr = np.zeros((2,2))
it = np.nditer([arr1, arr2, arr], [], [['readonly'], ['readonly'], ['writeonly']])
for a1, a2, a in it:
    a = a1 if -a1 < a2 else a2
print arr
print it.operands[2]

我在 arr 和 it.operands[2] 中得到的结果全为零，但我希望 arr1 或arr2。在迭代中为 arr 赋值的正确方法是什么？

Answer 1

在 Python 中执行 a = 将简单地重新绑定(bind)局部变量 a；它不会影响 a 包含的内容。

使用nditer，迭代变量a1、a2 和a 实际上是0 维数组。因此，要更改 a，请使用(有点奇怪的)a[()] = 语法:

for a1, a2, a in it:
    a[()] = a1 if -a1 < a2 else a2

不过请注意，使用 np.where 可以大大简化您的整个代码:

import numpy as np
arr1 = - np.random.random((2,2))
arr2 = np.random.random((2,2))
arr = np.where(-arr1 < arr2, arr1, arr2)