如何在 numpy ndarray 上使用循环和中断语句

Question

我有这个 a 的数据数组：

>>>a
array([[1., 2., 3.],
      [4., 5., 6.]], dtype=float32)

现在我想对数据做这样的事情：

>>>for n in range (1,4):
>>>    a += 2
>>>    print(a)
[[3. 4. 5.]
 [6. 7. 8.]]
[[ 5.  6.  7.]
 [ 8.  9. 10.]]
[[ 7.  8.  9.]
 [10. 11. 12.]]

这会给我a的最终结果：

>>> a
array([[ 7.,  8.,  9.],
       [10., 11., 12.]], dtype=float32)

如果我希望每个元素的迭代在值 > 8 时停止， 这可能会给我这样的最终结果：

>>> a
array([[7., 8., 7.],
      [8., 7., 8.]], dtype=float32)

怎么做？

谢谢！

Answer 1

使用np.where

代码

a = np.array([[1., 2., 3.],
      [4., 5., 6.]], dtype=np.float32)

for i in range(1, 4):
    a = np.where(a + 2 <= 8, a+2, a)

替代方案（需要 Python 3.8+）

您可以使用 Walrus 运算符来避免计算 a + 2 两次：

for i in range(1, 4):
    a = np.where((p:=a + 2) <= 8, p, a)

输出

print(a)
array([[7., 8., 7.],
       [8., 7., 8.]], dtype=float32)

说明

Syntax :numpy.where(condition[, x, y])
Parameters:
condition : When True, yield x, otherwise yield y.

Answer 2

一种方法是：

import numpy as np

a = np.array([[1., 2., 3.],
              [4., 5., 6.]], dtype=np.float32)

for n in range(1, 4):
    a += 2 * ((a + 2) <= 8)

print(a)

输出

[[7. 8. 7.]
 [8. 7. 8.]]

这个想法是将 2 乘以一个布尔掩码，如果 2 可以添加到 a 的元素中，则该掩码为 1，否则为 0。

作为替代方案，您可以通过执行以下操作完全跳过 for 循环：

a = 8 - (a % 2)
print(a)

输出

[[7. 8. 7.]
 [8. 7. 8.]]

上述解决方案是基于a中的偶数最终为8，奇数为7的事实，这当然是假设a中的所有数字都小于8 。