如何通过二维矩阵A、B运算得到指定维度的三维矩阵？

现有两个二维矩阵：

• 矩阵A：维度为 (n_m, n_u)
• 矩阵B：维度为 (n_m, n)

需要生成维度为 (n, n_m, n_u) 的三维矩阵C，生成规则为：

• 取B的第1列，与A的每一列做元素级乘法，得到C的第1个二维切片（即C[0, :, :]）
• 取B的第2列，重复上述操作得到C的第2个二维切片（C[1, :, :]）
• 以此类推，处理完B的所有列

示例演示 #

输入矩阵 #

A = |1 2 3|
    |4 5 6|

B = |1 2|
    |3 4|

输出三维矩阵C #

C的结构为两个二维切片：

切片1（对应B的第1列）：
|1   2  3|
|12 15 18|

切片2（对应B的第2列）：
|2   4  6|
|16 20 24|

整体三维维度为 (2, 2, 3)（对应n=2, n_m=2, n_u=3）

代码实现（基于NumPy） #

利用NumPy的广播机制可以高效实现该需求，无需循环：

import numpy as np

# 定义示例矩阵
A = np.array([[1, 2, 3],
              [4, 5, 6]])
B = np.array([[1, 2],
              [3, 4]])

# 生成三维矩阵C
# 将B转置后扩展维度为(n, n_m, 1)，A扩展维度为(1, n_m, n_u)，利用广播相乘
C = B.T[:, :, np.newaxis] * A[np.newaxis, :, :]

# 输出结果
print("三维矩阵C：")
print(C)

输出结果 #

三维矩阵C：
[[[ 1  2  3]
  [12 15 18]]

 [[ 2  4  6]
  [16 20 24]]]

原理说明 #

• B.T[:, :, np.newaxis]：将B转置为(n, n_m)后，添加一个新维度，变成(n, n_m, 1)
• A[np.newaxis, :, :]：给A添加一个新维度，变成(1, n_m, n_u)
• 两者相乘时，NumPy的广播机制会自动将维度扩展匹配，最终得到(n, n_m, n_u)的三维矩阵，正好符合需求

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