CODEWITHSUNDEEP
pythonarraysnumpy
noSkill
noSkill

Reputation: 68

Numpy indexing with argsort

For 1d indexing array with the result of argsort on that array gives you a sorted array.

ar = numpy.random.random(4)
s = numpy.argsort(ar)
ar[s].shape
(4,)

But for 3d i am getting:

ar = numpy.random.random((2,3,2))
s = numpy.argsort(ar)
ar[s].shape
(2, 3, 2, 3, 2)

How can i get same behavior for 3d array as for 1d? Thanks!

Upvotes: 1

Views: 1268

Answers (1)

chw21
chw21

Reputation: 8140

numpy.argsort sorts a multi-dimensional array along its last axis, unless specified otherwise. Your s.shape is the same as ar.shape, but the fact that you can even use ar[s] without getting IndexErrors is just because you chose a nice shape to begin with.

You first have to think about what you really want to sort. Say you have:

[[8, 7],
 [6, 5],
 [4, 3]]

What do you want to get? Left to right:

[[7, 8],
 [5, 6],
 [3, 4]]

or top to bottom:

[[4, 3],
 [6, 5],
 [8, 7]]

or completely:

[[3, 4],
 [5, 6],
 [7, 8]]

The last one is probably easiest to flatten and then reshape.

arf = ar.flatten()
s = numpy.argsort(arf)
arf[s].reshape(ar.shape)

The first one is a bit harder:

s = numpy.argsort(ar)
numpy.array([[ar2[s2] for ar2, s2 in zip(ar1, s1)] for ar1, s1 in zip(ar, s)])

The last one is homework.

Upvotes: 3

Related Questions