獲取數組的索引數組

Question

如果我有這樣的多維數組：

a = np.array([[9,9,9],[9,0,9],[9,9,9]])

我想得到該數組中每個索引的數組，如下所示：

i = np.array([[0,0],[0,1],[0,2],[1,0],[1,1],...])

我發現這樣做的一種方法就是這樣，使用np.indices ：

i = np.transpose(np.indices(a.shape)).reshape(a.shape[0] * a.shape[1], 2)

但這似乎有點笨拙，特別是考慮到np.nonzero的存在幾乎可以實現我想要的。

是否有一個內置的numpy函數，它將生成2D numpy數組中每個項的索引數組？

Answer 1

這是一種更簡潔的方法（如果訂單不重要）：

In [56]: np.indices(a.shape).T.reshape(a.size, 2)
Out[56]: 
array([[0, 0],
       [1, 0],
       [2, 0],
       [0, 1],
       [1, 1],
       [2, 1],
       [0, 2],
       [1, 2],
       [2, 2]])

如果您想按預期順序使用，可以使用dstack ：

In [46]: np.dstack(np.indices(a.shape)).reshape(a.size, 2)
Out[46]: 
array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 0],
       [1, 1],
       [1, 2],
       [2, 0],
       [2, 1],
       [2, 2]])

對於第一種方法，如果您不想使用reshape另一種方法是使用np.concatenate()沿第一軸連接。

np.concatenate(np.indices(a.shape).T)