python - 如何根据条件计算numpy数组的每个元素

Question

我有一个包含多个元素的单列 numpy 数组。我想编写一个函数来获取每个元素并评估它是否满足条件。然后根据结果，它应该使用方程式来计算一些东西。计算结果应与输入数组大小相同。

这是我想做的一个简单例子。 (实际代码会更复杂)。我知道为什么它不起作用，但似乎找不到解决方案。

import numpy as np
array1 = np.arange(1,11,1)

def test(array1):
    value1 = 20
    i = 0
    value3 = array1[i]
    while array1 > value3 and i < value1:
        i =+ 1
        value3 = array1[i]

test(array1)

我试图找到一个解决方案:

• 我查看了 any() 和 all()，但这不是我需要的，因为我想分别查看每个元素。
• 我查看了 numpy.where()，但它只返回满足条件的元素，因此改变了数组大小，这不是我想要的。

编辑: 对于我的完整解决方案，Zulfiqaar 的 for 循环版本和 Siva Kowuru 显示的 where 函数都是可能的。由于 where 函数看起来更快更方便，我将其标记为已接受的答案。

Answer 1

您可以在 numpy.where(condition[, x, y]) 上添加可选参数以保持相同的数组大小。

当为True时，yield x，否则yield y

https://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html

---

与@Zulfiqaar 的答案相比，解决方案是

np.where(array1 < 3, array1 ** 2, array1 / 3)

请注意，条件只是一个 bool 数组，因此您可以使用 binary operators 一次组合多个条件。 .例如与另一个数组进行比较:

array2 = np.random.randint(1, 11, 10)

>>> array2
Out[]: array([ 6,  1,  6,  2, 10,  4,  9, 10,  3,  3])

>>> np.where((array1 < 3) & (array1 > 5), array1 ** 2, array1 / 3)
Out[]:
array([ 0.33333333,  4.        ,  1.        ,  1.33333333,  1.66666667,
        2.        ,  2.33333333,  2.66666667,  3.        ,  3.33333333])