Question

我有任意数量的A ndarray N。我想创建元组（数组或列表）的数组B，其中每个元组中的第一个N元素是索引，最后一个元素是A中该索引的值。

例如：

A = array([[1, 2, 3], [4, 5, 6]])

然后

B = [(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]

在没有for循环的情况下，在NumPy中执行此操作的最佳/最快方法是什么？

Answer 1

如果你有Python 3，那么（使用np.ndenumerate）非常简单（并且速度适中）：

>>> import numpy as np
>>> A = np.array([[1, 2, 3], [4, 5, 6]])
>>> [(*idx, val) for idx, val in np.ndenumerate(A)]
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]

如果你希望它适用于Python 3和Python 2，那将会有所不同，因为Python 2不允许在元组文字中进行迭代解包。但你可以使用元组连接（另外）：

>>> [idx + (val,) for idx, val in np.ndenumerate(A)]
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]

如果您想完全留在NumPy中，最好使用np.mgrid创建索引：

>>> grid = np.mgrid[:A.shape[0], :A.shape[1]]  # indices!
>>> np.stack([grid[0], grid[1], A]).reshape(3, -1).T
array([[0, 0, 1],
       [0, 1, 2],
       [0, 2, 3],
       [1, 0, 4],
       [1, 1, 5],
       [1, 2, 6]])

然而，这需要一个循环将其转换为元组列表......但是将它转换为列表列表会很容易：

>>> np.stack([grid[0], grid[1], A]).reshape(3, -1).T.tolist()
[[0, 0, 1], [0, 1, 2], [0, 2, 3], [1, 0, 4], [1, 1, 5], [1, 2, 6]]

如果没有可见 for，也可以使用元组列表 - 循环：

>>> list(map(tuple, np.stack([grid[0], grid[1], A]).reshape(3, -1).T.tolist()))
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]

即使没有可见的for - 循环tolist，list，tuple和map也会隐藏for - 循环在Python层。

对于任意维数组，您需要稍微改变后一种方法：

coords = tuple(map(slice, A.shape))
grid = np.mgrid[coords]

# array version
np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T
# list of list version
np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T.tolist()
# list of tuple version
list(map(tuple, np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T.tolist()))

ndenumerate方法适用于任何维度的数组而不会发生变化，根据我的时间，只会慢2-3倍。

Answer 2

您也可以使用np.ndindex执行此操作，尽管@ Mseifert的方法在时间和简洁性方面是无与伦比的。这里唯一的循环是用实际值压缩坐标生成器。（与其他答案相同。）

def tuple_index(a):
    indices = np.ndindex(*a.shape)
    return [(*i, j) for i, j in zip(indices, a.flatten())]

print(tuple_index(a))
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]

在NumPy中获取ndarray的索引和值

2 个答案: