如何无循环实现矩阵A与矩阵B每行的点积生成多维数组？

问题描述 #

给定两个矩阵：

# 2x1矩阵A
A = np.array([[2],
              [3]])

# 3x3矩阵B
B = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

需要构造一个多维数组，结构为[A·B(row1), A·B(row2), A·B(row3)]。当前通过循环实现的代码如下：

import numpy as np

# 2x1矩阵A
A = np.array([[2],
              [3]])

# 3x3矩阵B
B = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])
C = []
for b_row in B:
    C.append(A.dot([b_row]))
C = np.array(C)

print(C)

运行输出：

[[[ 2  4  6]
  [ 3  6  9]]

 [[ 8 10 12]
  [12 15 18]]

 [[14 16 18]
  [21 24 27]]]

需求：不使用循环，兼容任意列数的矩阵B完成该操作。

无循环实现方案 #

利用numpy的广播机制或矩阵维度调整即可实现，以下是两种高效方案：

方案一：广播元素乘法 #

通过扩展维度让A和B的形状匹配，利用广播完成逐行的元素乘法，直接得到目标结构：

import numpy as np

A = np.array([[2], [3]])
B = np.array([[1,2,3], [4,5,6], [7,8,9]])

# 将A扩展为(2,1)，B扩展为(1,3,3)，广播后相乘得到目标形状
C = A[:, np.newaxis] * B[np.newaxis, :, :]
# 等价简洁写法：C = A[:, None] * B[None, :, :]

方案二：矩阵乘法维度适配 #

调整B的维度后与A的转置做矩阵乘法，直接输出目标形状：

import numpy as np

A = np.array([[2], [3]])
B = np.array([[1,2,3], [4,5,6], [7,8,9]])

# 将B的每行转为(1,3)维度，与A.T(1,2)做矩阵乘法，得到(3,2,3)的目标结构
C = A.T @ B[:, None, :]

两种方案的输出结果与循环实现完全一致，且支持任意行数、列数的二维矩阵B。

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