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

我有这个形状的多维数组（500000,3,2,3），我们称之为数据。数据基本上是500000套3点，这3点中的每一个都分成其x和y坐标（因此为2）。形状中的最后3个代表3个点的不同旋转。现在，我得到了这个由0至2之间的500000个数字组成的1d数组，它告诉我要保留哪些旋转，我们称之为rot_index。我想构造一个多维形状的形状（500000,3,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]]