如何使用numpy.einsum添加冗余索引构造特定四维数组？

下面提供几种无需循环的高效实现方式：

方法1：利用索引广播赋值 #

直接初始化全零数组，再通过索引定位到Y[i,j,k,k]的位置赋值，这是最直观的方式：

import numpy as np

# 假设X是输入的N×N×N数组
N = X.shape[0]
Y = np.zeros((N, N, N, N))
# 利用np.arange生成匹配的索引，实现广播赋值
Y[:, :, np.arange(N), np.arange(N)] = X

方法2：结合单位矩阵的广播乘法 #

借助单位矩阵np.eye(N)的特性（仅对角线元素为1，其余为0），通过维度扩展实现广播相乘：

import numpy as np

N = X.shape[0]
# 将X扩展为N×N×N×1的数组，和N×N的单位矩阵广播相乘
Y = X[..., np.newaxis] * np.eye(N)

这里X[..., np.newaxis]给X增加了一个长度为1的最后一维，np.eye(N)会自动广播为1×1×N×N的维度，两者相乘后，仅当最后两维下标相等时结果为X的对应元素，其余为0，完全符合需求。

方法3：使用np.einsum结合单位矩阵 #

如果偏好使用einsum，可以通过引入单位矩阵避免重复输出下标的问题：

import numpy as np

N = X.shape[0]
eye = np.eye(N)
# 利用einsum实现X和单位矩阵的张量积，仅保留对应对角线的组合
Y = np.einsum('ijk,kl->ijkl', X, eye)

这里einsum的下标ijk,kl->ijkl表示将X的第三维和单位矩阵的第一维收缩，最终结果Y[i,j,k,l] = X[i,j,k] * eye[k,l]，完美匹配需求。

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