获取数组的索引数组

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)