使用二维数组的排序索引对 3D 数组进行排序

Question

我有 2D 和 3D numpy 阵列。 二维阵列A具有形状(N, 3) ，3d 阵列B具有形状(N, 3, 3) 。 我想A沿axis=1排序，然后将相同的排序应用于数组B沿axis=2排序。

我知道我能做到

sort_idxs = np.argsort(A, axis=1)

但后来我不知道如何以我需要排列B的方式应用sort_idxs 。 sort_idxs的形状为(N, 3)如A 。 不知何故，我需要将 sort_idxs 的第一维sort_idxs到B的第一维， map 将sort_idxs的第二维到B的第三维，并忽略B的第二维。 我怎样才能做到这一点？

Answer 1

这可以使用解决

sort_idxs = np.argsort(A, axis=1)
B_sorted = np.take_along_axis(B, sort_idxs[:, np.newaxis, :], axis=2)