如何使用Numpy将数组小于5的元素求平方、大于5的元素求立方

实现方案 #

你可以用np.where一行完成需求，也可以用布尔索引实现，两种方式性能差异不大，按需选择即可。

方法1：使用np.where（代码最简洁）

import numpy as np

# 给定原始数组
arr = np.array([[0, 2, 4, 6],
               [7, 9, 11, 13],
               [14, 16, 18, 20],
               [21, 23, 25, 27]],
               dtype=np.uint64)

# 核心运算：小于5平方，大于5立方，等于5保留原值
result = np.where(arr < 5, arr ** 2, np.where(arr > 5, arr ** 3, arr))

方法2：使用布尔索引（可读性更高）

result = arr.copy()
result[arr < 5] **= 2
result[arr > 5] **= 3

运算结果

array([[    0,     4,    16,   216],
       [  343,   729,  1331,  2197],
       [ 2744,  4096,  5832,  8000],
       [ 9261, 12167, 15625, 19683]], dtype=uint64)

说明

• 两种方案都是逐元素运算，完全兼容你给定的uint64数据类型
• 如果你需要修改等于5的元素的处理逻辑，调整对应分支的运算规则即可
• 数据量较大时两种方案性能基本一致，不需要纠结选择哪一种

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