NumPy中如何逐行对(m,2)数组应用自定义函数得到(m,1)数组

云原生阿岚

2026-8-29

实现方法

你遇到的核心问题是默认参数下的numpy.vectorize会对数组做逐元素广播，不会把整行作为函数入参，以下三种方案都可以避开显式for循环实现需求：

方案1：修正numpy.vectorize参数，适配行输入逻辑
给vectorize传入signature参数指定输入输出维度，就可以让函数每次接收一整行作为入参。注意原始自定义函数里直接用x == [1., 1.]会返回布尔数组，无法直接用于if判断，需要加np.all()做整行匹配校验：

import numpy as np

def func(x):
    if np.all(x == [1., 1.]):
        return 0.
    if np.all(x == [-1., 1.]) or np.all(x == [-1., -1.]):
        return 1.
    if np.all(x == [1., -1.]):
        return 2.

# 声明按一维向量输入、标量输出的向量化函数
vec_func = np.vectorize(func, signature='(n)->()')
# 输入arr为形状(m,2)的原始数组，最后调整维度到(m,1)
result = vec_func(arr).reshape(-1, 1)

方案2：纯向量化布尔索引实现（性能最优）
从你的判断规则看，数组元素只有1.和-1.两种取值，直接通过列条件筛行做批量赋值，所有计算都在C层完成，处理大规模数据时效率远高于封装循环的方案：

import numpy as np

def process_arr(arr):
    res = np.zeros(arr.shape[0], dtype=np.float64)
    # 第一列为-1的行，不管第二列是1还是-1都返回1
    res[arr[:, 0] == -1.] = 1.
    # 第一列为1、第二列为-1的行返回2
    res[(arr[:, 0] == 1.) & (arr[:, 1] == -1.)] = 2.
    return res.reshape(-1, 1)

result = process_arr(arr)

方案3：使用numpy.apply_along_axis直接按行应用函数
这是NumPy专门为沿指定轴应用函数提供的接口，写法更简洁，不需要额外配置vectorize参数：
```
result = np.apply_along_axis(func, axis=1, arr=arr).reshape(-1, 1)
```