生成包含另一个NumPy数组索引的NumPy数组 [英] Generate NumPy array containing the indices of another NumPy array

问题描述

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

I'd like to generate a np.ndarray NumPy array for a given shape of another NumPy array. The former array should contain the corresponding indices for each cell of the latter array.

示例1

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

Let's say we have a = np.ones((3,)) which has a shape of (3,). I'd expect

[[0]
 [1]
 [2]]

因为a中有a[0]，a[1]和a[2]，可通过它们的索引0，1和2进行访问.

since there is a[0], a[1] and a[2] in a which can be accessed by their indices 0, 1 and 2.

示例2

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

For a shape of (3, 2) like b = np.ones((3, 2)) there is already very much to write. I'd expect

[[[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]，第三行.因此，我们在生成的数组中的匹配位置处得到[0 0]，[0 1]，[1 0]，[1 1]，[2 0]和[2 1].

since there are 6 cells in b which can be accessed by the corresponding indices b[0][0], b[0][1] for the first row, b[1][0], b[1][1] for the second row and b[2][0], b[2][1] for the third row. Therefore we get [0 0], [0 1], [1 0], [1 1], [2 0] and [2 1] at the matching positions in the generated array.

非常感谢您抽出宝贵的时间.让我知道我是否可以以任何方式澄清这个问题.

Thank you very much for taking the time. Let me know if I can clarify the question in any way.

推荐答案

使用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 returns an array of index grid where each subarray represents an axis:

np.indices((3, 2))

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

然后用np.stack转置数组，从不同的轴为每个元素的堆栈索引:

Then transpose the array with np.stack, stacking index for each element from different axis:

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

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

这篇关于生成包含另一个NumPy数组索引的NumPy数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！