python - 在 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

如果你有 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 - 循环 tolist、list、tuple 和 map 确实在 Python 层隐藏了一个 for 循环。

---

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

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