获取数组的索引数组

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)