如何用Numpy高效优雅地根据ndarray元素值生成同形异类型数组？

问题描述 #

我有一个形状为S的numpy.ndarray数组A，想要创建一个形状为S + (N,)（示例中N=5）的numpy.ndarray数组B。希望根据A的每个元素值，调用不同的函数生成B的对应元素，且所有函数的返回数组形状一致。

示例需求的嵌套循环实现：

import numpy as np

def func1():
    return np.zeros(5)

def func2():
    return np.ones(5)

def func(val):
    return func2() if val & 1 else func1()

A = np.random.randint(0, 100, (10, 10))
B = np.zeros((10, 10, 5))  # 修正原代码中B的初始化错误
for i in range(10):
    for j in range(10):
        B[i, j] = func(A[i, j])

我尝试过Numpy API参考中“函数式编程”章节的函数，但它们要么不支持逐元素操作，要么无法处理元素维度的变更。

附言：感谢@bill提供的思路，使用np.apply_along_axis可将func一次性应用到整行，但该方法效率仍有优化空间。

高效/优雅的实现方式 #

1. 完全向量化操作（推荐，效率最高） #

如果函数逻辑可以转化为基于数组的条件判断，优先使用Numpy的向量化索引——这是底层优化的操作，性能远超循环或apply类方法。

针对示例中的奇偶判断逻辑，实现如下：

import numpy as np

A = np.random.randint(0, 100, (10, 10))
# 初始化结果数组，形状为A的形状 + (5,)
B = np.zeros(A.shape + (5,), dtype=np.float64)
# 对A中所有奇数位置，赋值为全1数组
B[A % 2 == 1] = np.ones(5)

2. 向量化函数（适用于复杂逻辑） #

如果函数逻辑无法直接转化为向量化判断，可使用np.vectorize包装自定义函数，通过指定signature参数明确输入输出的形状，实现优雅的逐元素映射：

import numpy as np

def func(val):
    return np.ones(5) if val & 1 else np.zeros(5)

# 定义向量化函数，指定输入为标量，输出为长度5的数组
vec_func = np.vectorize(func, otypes=[np.float64], signature='()->(5)')
B = vec_func(A)

3. 方法对比 #

• 完全向量化操作：效率最高，是Numpy的原生优化方式，避免了Python层面的循环开销。
• np.vectorize：写法优雅，但底层仍为Python循环，性能略低于向量化操作，适合逻辑复杂的场景。
• np.apply_along_axis：本质是逐行循环，性能不如前两种方法，仅适合快速实现而非追求效率的场景。

内容的提问来源于stack exchange，提问作者efficiency over readability