通过二维索引对三维numpy数组进行排序

时间:2019-11-04 09:16:07

标签: python arrays numpy

对此我确实有麻烦。我有一个三维的numpy数组,我想通过二维索引数组对其进行重新排序。实际上,将以编程方式确定数组,并且三维数组可以是二维或四维的,但是为了简单起见,如果两个数组都是二维的,则这是理想的结果:

ph = np.array([[1,2,3], [3,2,1]])
ph_idx = np.array([[0,1,2], [2,1,0]])
for sub_dim_n, sub_dim_ph_idx in enumerate(ph_idx):
    ph[sub_dim_n] = ph[sub_dim_n][sub_dim_ph_idx]

这使ph数组变成:

array([[1, 2, 3],
       [1, 2, 3]])

我想要哪个。在相同的情况下,任何人都可以帮忙,但是我没有ph而是一个三维数组(psh),例如:

psh = np.array(
    [[[1,2,3]], 
     [[3,2,1]]]
)

希望很清楚,请问是否不清楚。预先感谢!

2 个答案:

答案 0 :(得分:2)

如果您想要得到一个ph.shape形状的数组,则可以简单地np.squeeze ph_ixs使形状匹配,然后使用它来索引ph

print(ph)
[[[1 2 3]]
 [[3 2 1]]]

print(ph_idx)
[[0 1 2]
 [2 1 0]]

np.take_along_axis(np.squeeze(ph), ph_idx, axis=-1)

array([[1, 2, 3],
       [1, 2, 3]])

答案 1 :(得分:0)

因此,这里的线索已经在有用的注释中了,但是为了完整起见,它就像使用np.take_along_axis和2d数组的广播版本一样简单:

psh = np.array(
    [[[1,2,3]], 
     [[3,2,1]]]
)
ph_idx = np.array(
    [[0,1,2], 
     [2,1,0]]
)
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

如果3d数组的dim1包含多个元素,则具有以下优点:

psh = np.array(
    [[[1,2,3],[4,5,6],[7,8,9]], 
     [[3,2,1],[6,5,4],[9,8,7]]]
)
ph_idx = np.array([[0,1,2], [2,1,0]])
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

给出

array([[[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]],

       [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]])