Sort 3D array using sort indices of a 2D array

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?

Jagerber48 · Accepted Answer · 2022-09-22 18:48:20Z

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)

answered Sep 22, 2022 at 18:48

Jagerber48

1,12514 silver badges24 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Mechanic Pig Over a year ago

B[np.arange(10)[:, None, None], np.arange(3)[None, :, None], sort_idxs[:, None]] also is OK.

Collectives™ on Stack Overflow

Sort 3D array using sort indices of a 2D array

1 Answer 1

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Related