如何在NumPy中将m×n矩阵与p×n×q矩阵按特定方式运算得到q×n×m×p数组？

要实现你描述的操作，完全可以借助NumPy的矢量化广播机制和维度变换来完成，无需手动循环填充数组。以下是具体步骤和代码示例：

核心思路 #

我们需要让X（m×n）和Y（p×n×q）的维度对齐，通过广播实现对应元素相加，最后调整维度到目标形状q×n×m×p：

1. 调整X的维度，使其能在q和p维度上广播；
2. 调整Y的维度，使其能在m维度上广播；
3. 利用NumPy自动广播完成元素相加，直接得到目标形状的数组。

代码实现 #

import numpy as np

# 定义示例维度和数组
m, n, p, q = 2, 3, 4, 5
X = np.random.rand(m, n)
Y = np.random.rand(p, n, q)

# 调整X的维度：(m, n) → (1, n, m, 1)
# 先转置X让n维度处于第二位，再添加两个新轴用于广播
X_reshaped = np.expand_dims(np.expand_dims(X.T, axis=0), axis=-1)

# 调整Y的维度：(p, n, q) → (q, n, p) → (q, n, 1, p)
# 转置Y让q维度到首位、p维度到末位，再添加新轴用于m维度的广播
Y_reshaped = np.expand_dims(Y.transpose(2, 1, 0), axis=2)

# 广播相加得到结果
result = X_reshaped + Y_reshaped

# 验证结果形状（应为(q, n, m, p)）
print(result.shape)  # 输出: (5, 3, 2, 4)

维度变化详解 #

• X的变换：原形状(m, n) → 转置为(n, m) → 添加轴后变为(1, n, m, 1)。这个形状允许X在第一个维度（对应q）和最后一个维度（对应p）上广播。
• Y的变换：原形状(p, n, q) → 转置为(q, n, p) → 添加轴后变为(q, n, 1, p)。这个形状允许Y在第三个维度（对应m）上广播。
• 广播相加：两个调整后的数组会自动扩展到相同形状(q, n, m, p)，对应位置的元素完成相加，最终得到目标数组。

正确性验证 #

你可以随机选取位置验证结果：

# 选取任意索引
q_idx, n_idx, m_idx, p_idx = 0, 1, 0, 2

# 对比计算结果
calc_value = X[m_idx, n_idx] + Y[p_idx, n_idx, q_idx]
result_value = result[q_idx, n_idx, m_idx, p_idx]

print(np.isclose(calc_value, result_value))  # 输出: True，说明结果正确

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