优化NumPy数组向量化变换：寻求无循环的高效实现方案

针对示例的高效向量化方案 #

你的原代码通过列表推导遍历每一列再拼接结果，我们可以直接利用NumPy的广播和索引操作完全消除循环，性能提升显著：

import numpy as np

def transformation(current_tuple):
    return np.array([[current_tuple[0], current_tuple[1]],
                    [0, 3.*current_tuple[0]]])

array_with_two_rows = np.random.randint(0, 10, size=(2, 6))

# 原循环实现
original_result = np.vstack([transformation(x) for x in array_with_two_rows.T])

# 向量化实现
row0, row1 = array_with_two_rows
n_cols = array_with_two_rows.shape[1]
# 初始化结果数组，对应6个2×2矩阵堆叠后的形状
vectorized_result = np.empty((n_cols * 2, 2))
# 填充每个变换矩阵的第一行：直接复用原数组的列转置结果
vectorized_result[::2] = array_with_two_rows.T
# 填充每个变换矩阵的第二行：第一列全0，第二列为3倍的row0元素
vectorized_result[1::2, 0] = 0
vectorized_result[1::2, 1] = 3 * row0

# 验证结果一致性
print(np.allclose(original_result, vectorized_result))  # 输出True

通用场景的向量化改造思路 #

由于你提到实际变换逻辑可能变更，唯一确定的是输入为一维数组/元组，输出为二维数组，这里提供两种通用改造方案：

方案1：将变换函数改造为向量化版本（性能最优） #

把函数改为接收二维数组（每一行对应一个输入的一维元组），返回三维数组（每个元素对应一个输出的二维矩阵），最后通过reshape拼接成目标形状：

def transformation_vec(input_arr):
    # input_arr: (k, d) 数组，k是输入个数，d是输入维度
    k = input_arr.shape[0]
    # 假设输出为2×2矩阵，可根据实际需求调整形状
    result = np.empty((k, 2, 2))
    # 按变换逻辑批量填充三维数组
    result[:, 0, 0] = input_arr[:, 0]
    result[:, 0, 1] = input_arr[:, 1]
    result[:, 1, 0] = 0
    result[:, 1, 1] = 3 * input_arr[:, 0]
    return result

# 调用方式
input_transposed = array_with_two_rows.T
vectorized_result = transformation_vec(input_transposed).reshape(-1, 2)

这种方式完全利用NumPy的向量化运算，后续无论变换逻辑如何修改，只要能映射到对三维数组的批量操作即可保持高效。

方案2：使用np.apply_along_axis（快速改造） #

如果暂时无法完全改造变换函数，可以用np.apply_along_axis替代列表推导，它内部做了循环优化，代码更简洁：

# 沿轴1（每一列）应用变换函数，结果自动堆叠后reshape
vectorized_result = np.apply_along_axis(transformation, axis=0, arr=array_with_two_rows).reshape(-1, 2)

注意：np.apply_along_axis并非真正的向量化，性能提升不如方案1，仅适合快速迭代代码的场景。

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