CODEWITHSUNDEEP
pythonarraysnumpy
Jagerber48
Jagerber48

Reputation: 965

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

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?

Upvotes: 0

Views: 99

Answers (1)

Jagerber48
Jagerber48

Reputation: 965

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)

Upvotes: 1

Related Questions