生成包含另一个 NumPy 数组索引的 NumPy 数组

Question

我想为另一个 NumPy 数组的给定形状生成一个 np.ndarray NumPy 数组。前一个数组应包含后一个数组的每个单元格的相应索引。

示例 1

假设我们有 a = np.ones((3,))，其形状为 (3,)。我希望

[[0]
 [1]
 [2]]

因为a中有a[0]、a[1]和a[2]，可以通过它们的索引0、1和2.

示例 2

像b = np.ones((3, 2))这样(3, 2)的形状，已经有很多东西要写了。我希望

[[[0 0]
  [0 1]]

 [[1 0]
  [1 1]]

 [[2 0]
  [2 1]]]

因为b中有6个单元格可以通过相应的索引b[0][0]访问，第一行b[0][1]，b[1][0]，b[1][1] 第二行，b[2][0]，b[2][1] 第三行。因此，我们在生成的数组中的匹配位置得到 [0 0]、[0 1]、[1 0]、[1 1]、[2 0] 和 [2 1]。

非常感谢您抽出宝贵时间。如果我能以任何方式澄清问题，请告诉我。

Answer 1

使用 np.indices 和 np.stack 的一种方法：

np.stack(np.indices((3,)), -1)

#array([[0],
#       [1],
#       [2]])

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])

np.indices returns 索引网格数组，其中每个子数组代表一个轴：

np.indices((3, 2))

#array([[[0, 0],
#        [1, 1],
#        [2, 2]],        
#       [[0, 1],
#        [0, 1],
#        [0, 1]]])

然后用np.stack转置数组，不同轴的每个元素的堆叠索引：

np.stack(np.indices((3,2)), -1)

#array([[[0, 0],
#        [0, 1]],
#       [[1, 0],
#        [1, 1]],
#       [[2, 0],
#        [2, 1]]])

生成包含另一个 NumPy 数组索引的 NumPy 数组

Generate NumPy array containing the indices of another NumPy array

python

arrays

numpy

indices