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]]

Answer 1

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.