如何基于行索引使用numpy的apply_along_axis执行差异化运算

核心问题说明 #

np.apply_along_axis 传入的回调函数默认只能拿到当前切片的元素值，无法直接获取行索引，且你示例代码里的x.index//2==0是逻辑错误，判断行索引是否为偶数应该用取模运算行索引 % 2 == 0。

基于apply_along_axis的实现方法 #

我们可以用闭包维护一个计数器来记录当前处理的行号，代码如下：

import numpy as np

sa = np.array(np.arange(4))
sa_changed = np.repeat(sa.reshape(1, len(sa)), repeats=2, axis=0)

# 构造带行索引计数的处理函数
def create_process_func():
    current_row = 0
    def process_row(x):
        nonlocal current_row
        # 偶数行加10，奇数行加50
        res = x + 10 if current_row % 2 == 0 else x + 50
        current_row += 1
        return res
    return process_row

# 注意要沿行维度处理需指定axis=1，你之前用的axis=0是按列切片
result = np.apply_along_axis(create_process_func(), axis=1, arr=sa_changed)
print(result)

输出结果：

[[10 11 12 13]
 [50 51 52 53]]

更高效的向量化替代方案 #

如果处理大数组，不推荐使用apply_along_axis（本质是循环，性能较低），直接用numpy广播特性实现更高效：

# 生成对应每行的增量数组
add_val = np.where(np.arange(sa_changed.shape[0]) % 2 == 0, 10, 50).reshape(-1, 1)
result = sa_changed + add_val

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