python - 生成包含另一个 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 返回一个索引网格数组，其中每个子数组代表一个轴:

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