如何正确使用numpy.vectorize处理数组的数组，以子数组为函数输入？

问题分析 #

你之前的写法存在两个核心问题：一是误用了operator.index（它的作用是将对象转换为整数索引，而非获取数组元素），二是np.vectorize默认会扁平化输入数组，遍历每个元素而非子数组。要让vectorize作用于每个子数组，需要通过signature参数指定输入输出的形状，明确告诉它要处理的单元是子数组而非单个元素。

正确实现方式 #

方法1：基于operator.itemgetter的简洁实现

先通过operator.itemgetter(0)定义获取子数组第0个元素的函数，再给np.vectorize指定signature参数，声明输入为一维子数组、输出为标量：

import numpy as np
import operator

# 定义获取子数组第0个元素的函数
get_first_element = operator.itemgetter(0)
# 用signature指定输入输出形状：输入是任意长度的一维数组，输出是标量
vec_get_first = np.vectorize(get_first_element, signature='(n)->()')

# 测试输入
input_arr = np.array([[4,3,2],[90,7,6]])
result = vec_get_first(input_arr)
print(result.tolist())  # 输出 [4, 90]

方法2：自定义函数扩展

如果需要对每个子数组执行更复杂的自定义逻辑，同样可以用signature参数适配：

import numpy as np

def custom_subarray_func(sub_arr):
    # 示例：返回子数组第一个元素的平方
    return sub_arr[0] ** 2

# 指定输入输出形状
vec_custom = np.vectorize(custom_subarray_func, signature='(n)->()')

input_arr = np.array([[4,3,2],[90,7,6]])
print(vec_custom(input_arr).tolist())  # 输出 [16, 8100]

原写法无效的原因 #

1. 第一种写法：vidx = np.vectorize(operator.index)结合functools.partial(vidx,0)，相当于把0固定为vidx的第一个参数，但vidx默认会遍历原数组的每个元素（如4、3、2等），最终执行operator.index(0, 4)这类不符合函数参数要求的操作，逻辑完全错误。
2. 第二种写法：vidx = np.vectorize(functools.partial(operator.index,0))，同样会让vectorize遍历原数组的单个元素，执行operator.index(0, 4)这类无效调用，既达不到处理子数组的目的，也不符合operator.index的功能逻辑。

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