带索引的python数组操作

nhaq1z21  于 2021-09-08  发布在  Java
关注(0)|答案(1)|浏览(505)

我尝试设计一个数组操作,如下所述:
如果数组中的元素小于0,则新值应设置为0。否则,新值应为原始值的平方。

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

另外,我想让这个操作可以采取任意的输入维度(2d数组,3d数组…)
我只能想到这个解决方案:

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

我已经知道如何完成负片部分,但我不知道如何处理正方形部分,有人能帮我吗?
谢谢!

a1o7rhls

a1o7rhls1#

只需使用numpy的 where 功能:

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

相关问题