在NumPy中获取ndarray的索引和值

Question

我有一个维度任意的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)]

在NumPy中，最好/最快的方法是什么，可以避免使用for循环？

- Sia Rezaei

Answer 1

如果您拥有Python 3，一种非常简单（并且速度适中）的方法是使用np.ndenumerate（使用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)]

尽管在Python层中没有可见的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()))

< p > ndenumerate 方法适用于任何维度的数组，而且根据我的计时只会慢2-3倍。

- MSeifert

---

你刚刚为我的知识库增添了一笔。 - piRSquared

@MSeifert 谢谢，但我正在寻找一些没有for循环的东西。因为这些矩阵可能会非常大，而Python的for循环速度相对较慢。你认为这是可能的吗？（附注：只需要支持Python3） - Sia Rezaei

@SiaRezaei 如果您想要一个元组列表，则不行。如果数组可以的话，那么请看一下我的回答的第二部分（mgrid方法）。 - MSeifert

@SiaRezaei 是的，可以不用 for 循环实现。但要小心，我认为这两种方法都不会很“快”。 - MSeifert

@MSeifert 谢谢，我应该如何更改代码以使其适用于任意维度的 A？ - Sia Rezaei

如果你正在处理任意维度的数组，请使用ndenumerate方法。这种方法最多会慢2-3倍，但适用于任何维度。而且它实际上是非常干净和易于理解的代码。 - MSeifert

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)]

- Brad Solomon