如何实现支持任意维度的Numpy数组操作：负数置0非负数求平方

方法1：单行简洁实现（推荐） #

直接用np.where即可一步完成，天然支持任意维度的numpy数组：

import numpy as np

# 1维示例
arr = np.array([-2, 5, 2, -1])
result = np.where(arr < 0, 0, arr ** 2)
print(result) # 输出 [ 0 25  4  0]

# 2维示例，验证多维度支持
arr_2d = np.array([[-1, 3, -5], [2, -4, 6]])
result_2d = np.where(arr_2d < 0, 0, arr_2d **2)
print(result_2d)
# 输出
# [[ 0  9  0]
#  [ 4  0 36]]

方法2：基于原有思路修改 #

原有代码报错的原因是：input[input >= 0]是筛选出的符合条件的子集（形状和完整数组不一致），右侧input*input是完整数组，两边形状不匹配才会报错。
你已经把负数全部设为0了，直接对整个数组做平方运算即可，0平方后还是0，完全符合需求：

arr = np.array([-2, 5, 2, -1])
arr[arr < 0] = 0
result = arr ** 2
print(result) # 输出 [ 0 25  4  0]

方法3：更精简的逻辑等价实现 #

可以用np.maximum先把所有小于0的数替换为0，再整体平方，代码更短：

result = np.maximum(arr, 0) ** 2

以上所有方法都兼容任意维度的numpy数组输入，不需要额外修改逻辑。

内容的提问来源于stack exchange，提问作者ethan