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

Question

I have 2D and a 3D numpy array. The 2d array A has shape (N, 3) and the 3d array B has shape (N, 3, 3) . I want to sort A along axis=1 and then apply that same sorting to array B sorting along axis=2 .

I know I can do

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

but then I don't know how to apply sort_idxs in the way I need to array B . sort_idxs has a shape of (N, 3) like A . Somehow I need to map the first dimension of sort_idxs to the first dimension of B , map the second dimension of sort_idxs to the 3rd dimension of B , and ignore the second dimension of B . How can I do this?

Answer 1

This can be solved using

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