numpy: applying argsort to an array

Question

The argsort() function returns a matrix of indices that can be used to index the original array so that the result would match the sort() result.

Is there a way to apply those indices? I have two arrays, one is the array used for obtaining the sort order, and another is some associated data.

I would like to compute assoc_data[array1.argsort()] but that doesn't seem to work.

Here's an example:

z=array([1,2,3,4,5,6,7])
z2=array([z,z*z-7])
i=z2.argsort()

z2=array([[ 1,  2,  3,  4,  5,  6,  7],
          [-6, -3,  2,  9, 18, 29, 42]])
i =array([[1, 1, 1, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 1, 1]])

I would like to apply i to z2 (or another array with associated data) but I'm not sure how to do so.

Which axis do you want to sort along?

user545424
– user545424

2012-06-28 23:36:02 +00:00
Commented Jun 28, 2012 at 23:36 — user545424
– user545424, Commented Jun 28, 2012 at 23:36

Bi Rico · Accepted Answer · 2012-06-28 22:42:49Z

10

This is probably overkill, but this will work in the nd case:

import numpy as np
axis = 0
index = list(np.ix_(*[np.arange(i) for i in z2.shape]))
index[axis] = z2.argsort(axis)
z2[index]

# Or if you only need the 3d case you can use np.ogrid.

axis = 0
index = np.ogrid[:z2.shape[0], :z2.shape[1], :z2.shape[2]]
index[axis] = z2.argsort(axis)
z2[index]

answered Jun 28, 2012 at 22:42

Bi Rico

25.9k3 gold badges57 silver badges75 bronze badges

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

1 Comment

Jason S Over a year ago

oddly enough I needed this again exactly a year later, and in searching for how to do it, I ran across the question I asked before.... I finally understand what this does. It doesn't look like overkill, btw.

Community · Accepted Answer · 2017-05-23 11:54:17Z

5

You're lucky I just got my masters degree in numpyology.

>>> def apply_argsort(a, axis=-1):
...     i = list(np.ogrid[[slice(x) for x in a.shape]])
...     i[axis] = a.argsort(axis)
...     return a[i]
... 
>>> a = np.array([[1,2,3,4,5,6,7],[-6,-3,2,9,18,29,42]])
>>> apply_argsort(a,0)
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

For an explanation of what's going on, see my answer to this question.

edited May 23, 2017 at 11:54

CommunityBot

11 silver badge

answered Jun 28, 2012 at 23:56

user545424

16.3k11 gold badges61 silver badges72 bronze badges

Comments

vozman · Accepted Answer · 2019-02-19 14:06:39Z

3

Use np.take_along_axis

np.take_along_axis(z2, i, axis=1)
Out[31]: 
array([[ 1,  2,  3,  4,  5,  6,  7],
       [-6, -3,  2,  9, 18, 29, 42]])

answered Feb 19, 2019 at 14:06

vozman

1,4251 gold badge18 silver badges25 bronze badges

1 Comment

Jason S Over a year ago

that's not the answer I am looking for, but it does look like np.take_along.axis can be used for this purpose, if called correctly.

Jason S · Accepted Answer · 2012-06-28 22:04:30Z

0

Aha, figured it out.

In [274]: z2[i,range(z2.shape[1])]
Out[274]:
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

answered Jun 28, 2012 at 22:04

Jason S

191k176 gold badges637 silver badges1k bronze badges

2 Comments

tiho Over a year ago

Plus you didn't say what i was ;)

tiho Over a year ago

Oops sorry, got confused by the fact there are two i's in your question and I was looking at the "wrong" one (the first one, that doesn't work when plugged in your answer)

Collectives™ on Stack Overflow

numpy: applying argsort to an array

4 Answers 4

1 Comment

Comments

1 Comment

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

1 Comment

Comments

1 Comment

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related