如何用Numpy实现将不同维度数组按元素组合生成新数组C

已知二维数组A（形状为(m, n)）和B（形状为(k, p)），需生成数组C，其中C包含A的每个子数组与B的每个子数组的拼接结果，最终C的形状为(m*k, n+p)。以下是两种高效的Numpy实现方案：

方案一：利用广播机制（内存效率更高） #

通过扩展数组维度实现广播，避免重复存储原始数组元素：

import numpy as np

# 示例输入
A = np.array([[1, 2, 3], [4, 5, 6]])
B = np.array([['a', 'b'], ['c', 'd'], ['e', 'f']])

# 扩展维度，使A和B可广播配对
A_expanded = A[:, np.newaxis, :]  # 形状变为 (2, 1, 3)
B_expanded = B[np.newaxis, :, :]  # 形状变为 (1, 3, 2)

# 在第三个维度拼接后展平为二维数组
C = np.concatenate([A_expanded, B_expanded], axis=2).reshape(-1, A.shape[1] + B.shape[1])

print(C)

输出结果：

[['1' '2' '3' 'a' 'b']
 ['1' '2' '3' 'c' 'd']
 ['1' '2' '3' 'e' 'f']
 ['4' '5' '6' 'a' 'b']
 ['4' '5' '6' 'c' 'd']
 ['4' '5' '6' 'e' 'f']]

原理说明 #

• 扩展维度后，A的每个行会自动与B的所有行配对，形成(m, k, n+p)的三维中间数组。
• reshape(-1, ...)自动计算第一维度的长度（即m*k），将三维数组展平为目标二维结构。

方案二：使用repeat与tile（逻辑更直观） #

通过重复A的行、平铺B的数组，再直接拼接：

import numpy as np

# 示例输入
A = np.array([[1, 2, 3], [4, 5, 6]])
B = np.array([['a', 'b'], ['c', 'd'], ['e', 'f']])

# 将A的每一行重复B的行数次数
A_repeated = np.repeat(A, B.shape[0], axis=0)
# 将B整体重复A的行数次数
B_tiled = np.tile(B, (A.shape[0], 1))

# 拼接得到目标数组C
C = np.concatenate([A_repeated, B_tiled], axis=1)

print(C)

输出结果与方案一完全一致。

原理说明 #

• np.repeat(A, k, axis=0)：把A的每个行连续重复k次（k是B的行数），得到形状为(m*k, n)的数组。
• np.tile(B, (m, 1))：把B在第一维度重复m次（m是A的行数），得到形状为(m*k, p)的数组。
• 最后在列维度拼接两个数组，直接得到目标C。

两种方案均为Numpy向量化操作，比Python循环效率高几个数量级，适合处理大规模数组。

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