在 NumPy 中获取 ndarray 的索引和值

Question

我有一个任意维数 N 的 ndarray A。我想创建一个数组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

您也可以使用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)]

Answer 2

如果您有 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-loop:

>>> 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-loop 不可见，tolist、list、tuple 和 map 也会在 Python 层中隐藏 for-loop。

---

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

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 倍。