Question

给定矩阵a a.shape == (M, N, O)，是否有更好的生成方式：

indices = (
    np.arange(M).reshape(M, 1, 1),
    np.arange(N).reshape(1, N, 1),
    np.arange(O).reshape(1, 1, O)
)

还适用于更高维度的数组？

我可以接近np.indices，但这会返回一个封闭的网格（所有条目的形状为M,N,O），而不是一个开放的网格。

Answer 1

您可以使用np.ix_执行此操作：

np.ix_(np.arange(M), np.arange(N), np.arange(O))

来自文档：

从多个序列构建一个开放网格。

此函数采用N个1-D序列并返回N个输出，每个N个维度，使得除一个维度外，形状为1，非单位形状值的维度循环遍历所有N个维度。

这会产生与indices此处（M, N, O = 2, 3, 4）：

相同的结果

(array([[[0]],

        [[1]]]), 

 array([[[0],
         [1],
         [2]]]), 

 array([[[0, 1, 2, 3]]]))

Answer 2

如果单独指定M，N和O，

np.ogrid可能是最简洁的方法：

M, N, O = 2, 3, 4

indices = np.ogrid[:M, :N, :O]

print(indices)
# [array([[[0]],

#        [[1]]]), array([[[0],
#         [1],
#         [2]]]), array([[[0, 1, 2, 3]]])]

如果输入只是数组本身，我能想到的最简洁的方法是：

np.ix_(*(np.r_[:s] for s in A.shape))

打开数组索引的网格

2 个答案: