Python numpy sort array according to another array and broadcast over an axis

Question

I have done this a million times before - sorting one array according to another. But this time it is just slightly more complicated and I have been stumped how to do it. Let me explain. I have two arrays, say A:

[[1.59956565 1.16421459]
 [1.21548342 1.63884363]
 [0.73023302 0.54681896]
 [2.02628432 1.32127994]
 [0.2132793  0.26559821]
 [0.38242608 0.30073228]]

and B:

[[[ 0.93634073  0.35109262]
  [-0.35109262  0.93634073]]

 [[-0.63561769  0.77200398]
  [ 0.77200398  0.63561769]]

 [[ 0.8331935   0.55298155]
  [-0.55298155  0.8331935 ]]

 [[ 0.96691332  0.25510513]
  [-0.25510513  0.96691332]]

 [[-0.41372983  0.91039971]
  [ 0.91039971  0.41372983]]

 [[ 0.84228545  0.53903174]
  [-0.53903174  0.84228545]]]

i.e., B's dimension is 1 more than A's. I want to fist sore A along the last axis:

[[1.16421459 1.59956565]
 [1.21548342 1.63884363]
 [0.54681896 0.73023302]
 [1.32127994 2.02628432]
 [0.2132793  0.26559821]
 [0.30073228 0.38242608]]

and then also sort the middle axis of B according to this sort, e.g. B should become:

[[[-0.35109262  0.93634073]
  [ 0.93634073  0.35109262]]

 [[-0.63561769  0.77200398]
  [ 0.77200398  0.63561769]]

 [[-0.55298155  0.8331935 ]
  [ 0.8331935   0.55298155]]

 [[-0.25510513  0.96691332]
  [ 0.96691332  0.25510513]]

 [[-0.41372983  0.91039971]
  [ 0.91039971  0.41372983]]

 [[-0.53903174  0.84228545]
  [ 0.84228545  0.53903174]]]

How can I do this with a view or a slice using the argsort of A? I tried but got nowhere, because there is one more axis in B.

It's unclear how your sorting of A should affect the sorting of B - it looks like B was simply sorted along one of its axes? — Grismar
– Grismar, Commented Jan 23, 2022 at 0:23
I have edited the question to make it more clear from the data. In B, its second-last dimension should be sorted according to the last dimension of A. — John
– John, Commented Jan 23, 2022 at 0:28
I'm sorry, but it is still unclear to me how your sorting of A is supposed to affect the sorting of B. Instead of providing data only, that you suppose makes clear how the sorting should be affected, can you point out exactly where and how the desired sorting of B is affected by that of A, and how it would change if you change a single value, for example? (it currently looks like B is sorted along axis 1, but only in the first 'column') — Grismar
– Grismar, Commented Jan 23, 2022 at 0:34
A has dimensions (n, 2). B has dimensions (n, 2, 2). Forget about the last dimension of B, it should stay fixed. In each 2-sized array in A, if the first element is > than the second, then switch the middle dimension of array B at that index. Does that make sense? — John
– John, Commented Jan 23, 2022 at 0:46
It does, and then it would seem @fsl's solution is what you're after? — Grismar
– Grismar, Commented Jan 23, 2022 at 1:33

fsl · Accepted Answer · 2022-01-23 00:58:25Z

3

Here is one way using take_along_axis:

np.take_along_axis(B, A.argsort(1)[:, :, None], axis=1)

Output:

array([[[-0.35109262,  0.93634073],
        [ 0.93634073,  0.35109262]],

       [[-0.63561769,  0.77200398],
        [ 0.77200398,  0.63561769]],

       [[-0.55298155,  0.8331935 ],
        [ 0.8331935 ,  0.55298155]],

       [[-0.25510513,  0.96691332],
        [ 0.96691332,  0.25510513]],

       [[-0.41372983,  0.91039971],
        [ 0.91039971,  0.41372983]],

       [[-0.53903174,  0.84228545],
        [ 0.84228545,  0.53903174]]])

answered Jan 23, 2022 at 0:58

fsl

3,2801 gold badge12 silver badges21 bronze badges

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

Collectives™ on Stack Overflow

Python numpy sort array according to another array and broadcast over an axis

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related