如何将形状为[N,4]与[M]的Numpy数组合并为[N*M,4]的数组？

我有两个Numpy数组：

• 形状为[N,4]的数组x1，示例如下：

import numpy as np
x1 = np.array([[1,2,3,4],
               [2,3,4,5],
               [7,3,2,1]])

• 形状为[M]的数组x2，示例如下：

x2 = np.array([1,2,3,4])

希望得到形状为[N*M,4]的最终数组x3，其中c为常数，元素需满足如下规则：

x3 = [[1+c,2+x2[0],3-c,4-x2[0]],
      [2+c,3+x2[0],4-c,5-x2[0]],
      [7+c,3+x2[0],2-c,1-x2[0]],
      [1+c,2+x2[1],3-c,4-x2[1]],
      [2+c,3+x2[1],4-c,5-x2[1]],
      [7+c,3+x2[1],2-c,1-x2[1]],
      [1+c,2+x2[2],3-c,4-x2[2]],
      [2+c,3+x2[2],4-c,5-x2[2]],
      [7+c,3+x2[2],2-c,1-x2[2]],
      [1+c,2+x2[3],3-c,4-x2[3]],
      [2+c,3+x2[3],4-c,5-x2[3]],
      [7+c,3+x2[3],2-c,1-x2[3]]]

当前可通过循环遍历x1的行实现，但想找到更简便的方法。

不用循环，借助Numpy的广播特性就能高效完成，步骤如下：

1. 分析变换规则 #

先拆解每一列的计算逻辑：

• 第0列：x1[:,0] + c（所有行都加常数c）
• 第1列：x1[:,1]分别与x2的每个元素相加
• 第2列：x1[:,2] - c（所有行都减常数c）
• 第3列：x1[:,3]分别与x2的每个元素相减

2. 扩展维度实现广播配对 #

通过扩展数组维度，让x1和x2满足广播条件：

• 将x1从(N,4)扩展为(N,1,4)，新增一个维度对应x2的元素数量
• 将x2从(M)扩展为(1,M,1)，新增两个维度分别对应x1的行和列位置

这样两者可以自动广播到(N,M,4)的形状，完成所有元素对的计算。

3. 完整代码实现 #

import numpy as np

# 示例输入
x1 = np.array([[1,2,3,4],
               [2,3,4,5],
               [7,3,2,1]])
x2 = np.array([1,2,3,4])
c = 1  # 自定义常数

# 扩展维度，满足广播条件
x1_expanded = x1[:, np.newaxis, :]  # shape: (3,1,4)
x2_expanded = x2[np.newaxis, :, np.newaxis]  # shape: (1,4,1)

# 按规则计算每一列
x3 = np.empty((x1.shape[0], x2.shape[0], 4))
x3[..., 0] = x1_expanded[..., 0] + c
x3[..., 1] = x1_expanded[..., 1] + x2_expanded
x3[..., 2] = x1_expanded[..., 2] - c
x3[..., 3] = x1_expanded[..., 3] - x2_expanded

# 展平为(N*M,4)的最终形状
x3 = x3.reshape(-1, 4)

print(x3)

4. 代码说明 #

• 扩展维度后，Numpy会自动完成x1每行与x2每个元素的配对计算，无需手动循环
• reshape(-1,4)会自动计算第一维度的长度（N*M），将三维数组展平为目标二维数组
• 这种向量化运算的效率远高于循环，尤其当N、M数值较大时优势更显著

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