向量化从具有索引列表的多维数组中提取子多维数组

如何解决向量化从具有索引列表的多维数组中提取子多维数组

我有这个形状为多维的数组（500000,3,2,3），我们称之为数据。数据基本上是500000套3点，这3点中的每一个都分成其x和y坐标（因此为2）。形状的最后3个代表3个点的不同旋转。现在，我得到了这个由0至2之间的500000个数字组成的1d数组，它告诉我要保留哪些旋转，我们称之为rot_index。我想构造一个多维形状的形状（500000,2），仅保留正确旋转的数据点。关于如何从原始数据数组中提取具有正确索引的数据的任何想法？我尝试过类似的方法，但是没有用

data[:,:,rot_index]

编辑：

这是一些示例数据（给出10组点而不是500000点）

data = 
[[[[0.70846822 0.98552876 0.66736535]
   [0.         0.         0.        ]]

  [[0.66736535 0.70846822 0.98552876]
   [1.54545219 2.39798549 2.33974762]]

  [[0.98552876 0.66736535 0.70846822]
   [3.88519982 3.94343768 4.73773311]]]


 [[[0.8132551  1.18845796 1.53004225]
   [0.         0.         0.        ]]

  [[1.18845796 1.53004225 0.8132551 ]
   [1.43211754 2.58720625 2.26386152]]

  [[1.53004225 0.8132551  1.18845796]
   [4.01932379 4.85106777 3.69597906]]]


 [[[0.66123513 0.93651048 0.83170562]
   [0.         0.         0.        ]]

  [[0.93651048 0.83170562 0.66123513]
   [2.09747072 2.38383457 1.80188002]]

  [[0.83170562 0.66123513 0.93651048]
   [4.48130529 4.18571459 3.89935074]]]


 [[[1.31047414 0.67740955 1.42020073]
   [0.         0.         0.        ]]

  [[0.67740955 1.42020073 1.31047414]
   [1.66061575 1.97600777 2.64656179]]

  [[1.42020073 1.31047414 0.67740955]
   [3.63662352 4.62256956 4.30717753]]]


 [[[1.4085555  1.64177102 0.27708893]
   [0.         0.         0.        ]]

  [[0.27708893 1.4085555  1.64177102]
   [0.62154257 3.04315813 2.61848461]]

  [[1.64177102 0.27708893 1.4085555 ]
   [3.24002718 3.6647007  5.66164274]]]


 [[[0.48080385 0.85910831 0.52342904]
   [0.         0.         0.        ]]

  [[0.52342904 0.48080385 0.85910831]
   [1.08970318 2.57102289 2.62245924]]

  [[0.85910831 0.52342904 0.48080385]
   [3.71216242 3.66072607 5.19348213]]]


 [[[1.13610207 1.51237019 0.47256909]
   [0.         0.         0.        ]]

  [[1.51237019 0.47256909 1.13610207]
   [2.92304081 2.59328103 0.76686347]]

  [[0.47256909 1.13610207 1.51237019]
   [5.51632184 3.3601445  3.68990428]]]


 [[[1.08397801 1.16506242 0.84703646]
   [0.         0.         0.        ]]

  [[1.16506242 0.84703646 1.08397801]
   [2.37250664 2.04419242 1.86648625]]

  [[0.84703646 1.08397801 1.16506242]
   [4.41669906 3.91067866 4.23899289]]]


 [[[0.98734317 1.11177984 0.90283297]
   [0.         0.         0.        ]]

  [[1.11177984 0.90283297 0.98734317]
   [2.25981006 2.13666143 1.88671382]]

  [[0.90283297 0.98734317 1.11177984]
   [4.39647149 4.02337525 4.14652387]]]


 [[[1.94118244 1.14738719 1.98251535]
   [0.         0.         0.        ]]

  [[1.14738719 1.98251535 1.94118244]
   [1.83291888 1.90183408 2.54843234]]

  [[1.98251535 1.94118244 1.14738719]
   [3.73475296 4.45026642 4.38135123]]]]

这是我要保留的索引列表：

rot_index = np.array([1 2 1 1 1 1 1 2 1 1])

因此，例如，如果您考虑

data[0,0] = [[0.70846822 0.]
 [0.66736535 1.54545219]
 [0.98552876 3.88519982]]
data[0,1] = [[0.98552876 0.]
 [0.70846822 2.39798549]
 [0.66736535 3.94343768]]
data[0,2] = [[0.66736535 0.]
 [0.98552876 2.33974762]
 [0.70846822 4.73773311]]

这些是同一样本的3个不同的“旋转”，如果我们看rot_index的第一个元素，它是1。所以我只想保留

data[0,1] = [[0.98552876 0.]
 [0.70846822 2.39798549]
 [0.66736535 3.94343768]]

解决方法

使用numpy advanced indexing，在此之下，combining advanced and basic indexing的特定子主题应该可以工作（其中data_array是具有data的numpy ndarray）：

result = data_array[range(500000),...,rot_index]

对于您的样本数据，将产生：

[[[0.98552876 0.        ]
  [0.70846822 2.39798549]
  [0.66736535 3.94343768]]

 [[1.53004225 0.        ]
  [0.8132551  2.26386152]
  [1.18845796 3.69597906]]

 [[0.93651048 0.        ]
  [0.83170562 2.38383457]
  [0.66123513 4.18571459]]

 [[0.67740955 0.        ]
  [1.42020073 1.97600777]
  [1.31047414 4.62256956]]

 [[1.64177102 0.        ]
  [1.4085555  3.04315813]
  [0.27708893 3.6647007 ]]

 [[0.85910831 0.        ]
  [0.48080385 2.57102289]
  [0.52342904 3.66072607]]

 [[1.51237019 0.        ]
  [0.47256909 2.59328103]
  [1.13610207 3.3601445 ]]

 [[0.84703646 0.        ]
  [1.08397801 1.86648625]
  [1.16506242 4.23899289]]

 [[1.11177984 0.        ]
  [0.90283297 2.13666143]
  [0.98734317 4.02337525]]

 [[1.14738719 0.        ]
  [1.98251535 1.90183408]
  [1.94118244 4.45026642]]]