如何对Numpy数组中符合特定索引条件的元素应用函数

我有一个任意尺寸的3D NumPy数组，以及一个函数fun()，该函数接收向量[x1,x2,...,xn]并返回每个元素平方后的向量，即[x1**2,x2**2,...,xn**2]。我希望将该函数应用于索引符合特定条件的元素，例如对索引形式为[0,0],[0,2],...,[2,0],[2,2],...（两个坐标均能被2整除）的每个元素应用fun()。我知道可以用array[array%2 == 0]判断值是否能被2整除，但不知道如何对索引进行条件判断。我考虑过使用apply_over_axis并在函数内检查坐标是否符合要求，但认为这种方法效率太低，请问有没有高效的实现方式？

方法一：利用NumPy切片（规则索引场景） #

如果你的索引条件是前两个维度的索引均为偶数（如示例中的规则模式），直接使用切片操作是最高效的方案——切片是NumPy原生优化的向量化操作，完全避开Python循环开销：

import numpy as np

# 示例3D数组
arr = np.random.randint(0, 10, (4, 4, 3))

# 定义目标函数（此处为平方，可替换为任意逐元素处理函数）
def fun(vec):
    return vec ** 2

# 选中前两个维度索引为偶数的所有元素，应用函数并赋值
arr[::2, ::2, :] = fun(arr[::2, ::2, :])

• 解释：::2表示从索引0开始以步长2选取元素，arr[::2, ::2, :]精准定位了所有(i,j,k)中i、j为偶数的元素，第三个维度的:表示选取该维度的全部元素。

方法二：布尔掩码（通用索引场景） #

如果索引条件更复杂（比如不规则的索引组合），可以通过生成布尔掩码来精准筛选目标元素，同样是向量化操作，效率远高于apply_over_axis：

import numpy as np

arr = np.random.randint(0, 10, (4, 4, 3))

def fun(vec):
    return vec ** 2

# 生成各维度的索引网格
i, j, k = np.indices(arr.shape)
# 创建掩码：前两个维度索引均为偶数
mask = (i % 2 == 0) & (j % 2 == 0)

# 对掩码选中的元素应用函数
arr[mask] = fun(arr[mask])

• 解释：np.indices(arr.shape)生成对应每个维度的索引数组，通过逻辑运算得到布尔掩码，最后用掩码索引直接定位目标元素并完成函数应用。

效率对比说明 #

apply_over_axis本质是在Python层面对每个元素/轴做循环处理，无法利用NumPy的C级向量化优化，速度极慢；而切片和掩码都是NumPy底层实现的向量化操作，所有计算在C语言层面完成，效率提升几个数量级。

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