numpy: reorder array by specified values

Question

I have a matrix:

A = [ [1,2],
      [3,4],
      [5,6] ]

and a vector of values:

V = [4,6,2]

I would like to reorder A by 2nd column, using values from V. The result should be:

A = [ [3,4],
      [5,6], 
      [1,2] ] # 2nd columns' values have the same order as V

How to do it?

Answer 1

First, we need to find the indicies of the values in the second column of A that we'd need to match the order of V. In this case, that's [1,2,0]. Once we have those, we can just use numpy's "fancy" indexing to do the rest.

So, you might do something like this:

import numpy as np
A = np.arange(6).reshape((3,2)) + 1
V = [4,6,2]
column = A[:,1].tolist()
order = [column.index(item) for item in V]
print A[order,:]

If you want to avoid python lists entirely, then you can do something like what's shown below. It's hackish, and there may be a better way, though...

We can abuse numpy.unique to do this... What I'm doing here is depending on a particular implementation detail (unique seems to start at the end of the array) which could change at any time... That's what makes it an ugly hack.

import numpy as np
A = np.arange(6).reshape((3,2)) + 1
V = np.array([4,6,2])
vals, order = np.unique(np.hstack((A[:,1],V)), return_inverse=True)
order = order[-V.size:]
print A[order,:]

Answer 2

@JoeKington's numpy solution is very clever, but it relies on A[:,1] being in sorted order. Here is a fix for the general case:

import numpy as np

np.random.seed(1)
N=5
A = np.arange(2*N).reshape((-1,2))+100
np.random.shuffle(A)
print(A)

If A looks like this:

[[104 105]
 [102 103]
 [108 109]
 [100 101]
 [106 107]]

and V

V = A[:,1].copy()
np.random.shuffle(V)
print(V)

looks like this:

[105 109 107 101 103]

then we use Joe's solution:

vals, order = np.unique(np.hstack((A[:,1],V)), return_inverse=True)

but save both the order of A[:,1] and V:

a_order = order[:V.size]
v_order = order[-V.size:]

and sort A (by forming A[np.argsort(a_order)]) before reordering with v_order:

print A[np.argsort(a_order)][v_order]

[[104 105]
 [108 109]
 [106 107]
 [100 101]
 [102 103]]

(A[np.argsort(a_order)] is A sorted according to its second column.)

---

Note that np.unique always returns the array in sorted order. The documentation guarantees with return_inverse=True that the returned indices are the indices of the unique array that reconstructs the original array. That is, if you call np.unique like this:

uniq_arr, indices = np.unique(arr, return_inverse=True)

you are guaranteed that

unique_arr[indices] = arr

Because you can rely on this relationship, Joe's method does not depend on a mere implementation detail -- unique will always behave this way. (Famous last words -- considering what happened to the order of output arguments returned by np.unique1d ... but never mind that :))