带索引的python数组操作

Question

我尝试设计一个数组操作，描述如下：

如果数组中的元素小于 0，则新值应设置为 0。 否则新值应该是原始值的平方。

e.g.
import numpy as np
input = np.array([-2, 5, 2, -1])
#some operation ...
#input should be: [0, 25, 4, 0]

另外，我想让这个操作可以任意输入维度（2D数组，3D数组...）

我只能想到这个解决方案：

input = np.array([-2, 5, 2, -1])
input[input < 0] = 0  # input = [0 5 2 0]
#input[input >= 0] = input*input   <- this will cause error

我已经知道如何完成否定部分，但我不知道如何处理方形部分，有人可以帮助我吗？

非常感谢！

Answer 1

只需使用 numpy 的 where 函数：

output = np.where(input <= 0, 0, input**2)