以索引列表和值列表的形式查看ndarray

Question

我有一个2或6维的ND数组。例如：

nd_array = [[90, 80],
[70, 60], 
[50, 40]]

我想把它看作一个索引列表，比如：

nd_array_transformed = [[90, 0, 0],
[80, 0, 1],
[70, 1, 0],
[60, 1, 1],
[50, 2, 0], 
[40, 2, 1]]

或类似的表示。

我为2D案例找到了这个代码，但我正在寻找一个更通用的解决方案：

np.array([[nd_array[row][col], row, col] for row in range(nd_array.shape[0]) 
for col in range(nd_array.shape[1])])

在一般情况下，可能使用*nd_array.shape进行这种转换吗？

Answer 1

您可以将itertools.product()与range()和map()一起使用以生成所有可能的索引，然后使用拆包运算符将索引添加到每个内部列表中。这里使用的例子是二维的，但与固定数量的for循环相反，使用开箱使这种方法可以扩展到任意数量的维度：

import numpy as np
from itertools import product
nd_array = np.array([[90, 80],
[70, 60],
[50, 40]])
result = np.array([[nd_array[indices], *indices]
for indices in product(*map(range, nd_array.shape))])
print(result)

---

用于构建我们在列表理解中迭代的内容的语法有点密集，因此这里有一个更详细的解释：

.shape产生一个元组，该元组给出数组的维度(例如(3, 2)(。所给出的例子都将采用这种形状。

我们希望能够生成由形状(例如(0, 0), (0, 1) ... (2, 1)(给出的所有可能的索引。我们使用map()和range()，取元组中的每个整数，并生成用该整数初始化的相应的range()(即(3, 2)被转换为包含range(3)和range(2)的可迭代对象(。然后(使用*(将该可迭代项解包为itertools.product()——函数调用本质上变为itertools.product(range(3), range(2))，生成一个笛卡尔乘积，为我们提供所有所需的索引。