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