对所有(x[k,i], y[j])配对求和生成指定shape numpy数组的高效方法

核心思路是利用numpy广播机制消除所有Python层面循环，仅需一行向量化代码即可实现需求：

import numpy as np

# 输入x.shape = (n, m), y.shape = (p,)
z = (x[:, np.newaxis, :] + y[np.newaxis, :, np.newaxis]).reshape(n, m * p)

实现原理 #

1. 维度扩展广播：给x新增中间维度得到形状为(n, 1, m)的数组，给y新增前后两个维度得到形状为(1, p, 1)的数组，两者广播相加后得到形状为(n, p, m)的中间数组，其中中间数组[k, j, i] = x[k, i] + y[j]，完全匹配配对求和逻辑。
2. 维度重塑：numpy默认采用C风格（行优先）的reshape规则，(n, p, m)的数组重塑为(n, m*p)后，原j, i对应的索引会自动转换为j * m + i，正好符合z[k, i + m * j]的赋值规则。
3. 内存开销说明：中间数组的总元素数为n*m*p，和最终输出数组的总元素数完全一致，不会产生额外内存占用。

正确性验证 #

可以用小规模样例验证和朴素算法结果完全一致：

# 测试样例
n, m, p = 2, 3, 2
x = np.arange(n*m).reshape(n, m)
y = np.arange(p)

# 向量化实现
z_vec = (x[:, np.newaxis, :] + y[np.newaxis, :, np.newaxis]).reshape(n, m*p)

# 朴素实现
z_naive = np.empty((n, m*p))
for k in range(n):
    for i in range(m):
        for j in range(p):
            z_naive[k, i + m * j] = x[k, i] + y[j]

print(np.allclose(z_vec, z_naive))  # 输出True

性能优势 #

该实现所有运算都在numpy底层C层面执行，没有Python循环的额外开销，即使n规模达到1e6，只要内存足够容纳最终数组，运行速度会比朴素三层Python循环快至少2~3个数量级，完全满足大规模数据的处理需求。

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