求基于einsum实现数组列行外积生成3D数组的更优方案

嘿，这个需求用np.einsum可以非常简洁高效地实现，一行代码就能搞定！先直接上核心实现，再给你拆解原理和验证方法：

核心代码 #

首先定义你的输入数组：

import numpy as np

A = np.arange(12).reshape(4, 3)
B = np.arange(-5, 7).reshape(3, 4)

然后用einsum一步生成目标3D数组C：

C = np.einsum('ik,kj->ijk', A, B)

下标含义拆解 #

这个einsum的下标规则正好匹配你的需求：

• 'ik'：对应A的形状(4,3)，其中i代表A的行维度（长度4），k代表A的列维度（长度3）
• 'kj'：对应B的形状(3,4)，其中k代表B的行维度（长度3），j代表B的列维度（长度4）
• 'ijk'：指定输出数组的维度顺序，最终形状是(4,4,3)——正好是你需要的：i（4）、j（4）对应输出的前两个维度，k（3）作为第三个维度存储每一组外积结果

验证结果正确性 #

我们可以直接验证每一个k对应的切片是否符合要求：

# 验证k=0的情况
assert np.allclose(C[:, :, 0], A[:, [0]] @ B[[0]])
# 验证k=1的情况
assert np.allclose(C[:, :, 1], A[:, [1]] @ B[[1]])
# 验证k=2的情况
assert np.allclose(C[:, :, 2], A[:, [2]] @ B[[2]])

如果没有抛出断言错误，说明结果完全符合你的要求。

为什么这个方案高效？ #

相比循环遍历每一个k计算外积，或者用广播手动调整维度，np.einsum的优势在于：

1. 代码可读性强：直接通过下标映射就能清晰表达维度之间的运算关系，不用纠结广播的维度扩展逻辑
2. 底层优化充分：NumPy内部对einsum做了大量优化，运算效率和直接用广播实现的高效方案相当，但代码更简洁

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