生成包含另一个NumPy数组的索引的NumPy数组

Question

我想为另一个NumPy数组的给定形状生成一个np.ndarrayNumPy数组。前一个数组应该包含后一个数组的每个单元格的相应索引。

例1

假设我们有a = np.ones((3,))，它的形状是(3,)。我想是的

[[0]
 [1]
 [2]]

由于a[0]、a[1]和{}可以通过它们的索引0、1和{}来访问。

例2

对于像(3, 2)这样的形状，已经有很多东西要写了。我想是的

^{pr2}$

由于b中有6个单元格，可以通过相应的索引b[0][0]，b[0][1]访问第一行，b[1][0]，b[1][1]用于第二行，b[2][0]，b[2][1]用于第三行。因此，我们得到[0 0]、[0 1]、[1 0]、[1 1]、[2 0]和{}。

非常感谢您抽出时间。如果我能以任何方式澄清这个问题，请告诉我。

Answer 1

使用np.indices和np.stack的一种方法：

np.stack(np.indices((3,)), -1)

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

np.stack(np.indices((3,2)), -1)

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

np.indices返回索引网格数组，其中每个子数组代表一个轴：

^{pr2}$

然后用np.stack对数组进行转置，对不同轴上的每个元素进行堆叠索引：

np.stack(np.indices((3,2)), -1)

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