Numpy: sort rows of an array by rows in another array

Question

I have a 2D array of "neighbors", and I want to re-order each row according to a corresponding row in another matrix (called "radii"). The below code works, but it uses a for loop over a numpy array, which I know is the incorrect way to do it. What is the correct numpy / broadcast solution to this re-ordering?

neighbors = np.array([[8,7,6], [3,2,1]])
radii = np.array([[0.4, 0.2, 0.1], [0.3, 0.9, 0.1]])

order = radii.argsort(axis=1)
for i in range(2):
    neighbors[i] = neighbors[i,order[i]]
print(neighbors)

# Result:
[[6 7 8]
 [1 3 2]]

Alex Riley · Accepted Answer · 2016-01-23 16:51:52Z

4

In NumPy you would write something like this:

>>> neighbors[np.arange(2)[:, None], order]
array([[6, 7, 8],
       [1, 3, 2]])

(More generally you'd write the first index as np.arange(order.shape[0])[:, None] instead.)

This works because np.arange(2)[:, None] looks like this:

array([[0],
       [1]])

and order looks like this:

array([[2, 1, 0],
       [2, 0, 1]])

For the fancy indexing, NumPy pairs off the arrays indexing each axis. The row index [0] is paired with the column index [2, 1, 0] and the new row is created in the order this determines. Similarly for [1] and [2, 0, 1] to determine the second row.

edited Jan 23, 2016 at 16:51

answered Jan 23, 2016 at 16:35

Alex Riley

178k46 gold badges274 silver badges247 bronze badges

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

1 Comment

Jason Maldonis Over a year ago

Thanks this is great! I think you have a small typo that threw me for a minute -- on the first line, I think your last ) should be a ]

Collectives™ on Stack Overflow

Numpy: sort rows of an array by rows in another array

1 Answer 1

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Related