How to sort a numpy matrix using a mask?

Question

I have two matrices A, B, Which look like this:

A = array([[2, 2, 1, 0, 8],
           [8, 2, 0, 3, 7],
           [3, 2, 6, 5, 3],
           [1, 4, 2, 5, 8],
           [2, 3, 7, 0, 3]])

B = array([[3, 7, 6, 8, 3],
           [0, 7, 4, 4, 3],
           [1, 2, 0, 0, 4],
           [8, 6, 6, 7, 1],
           [8, 1, 0, 4, 8]])

I am trying to sort A and B BUT I need B to be ordered with the mask from A.

I tried this:

mask = A.argsort()
A = A[mask]
B = B[mask]

However the return value is a shaped (5, 5, 5) matrix

The next snippet works, but is using two iterations. I need something faster. Has anybody an Idea ?

A = [row[order] for row, order in zip(A,mask)]
B = [row[order] for row, order in zip(B,mask)]

Answer 1

You can use fancy indexing. The result will be the same shape as your indices broadcasted together. Your column index is already the right shape. A row index of size (A.shape[0], 1) would broadcast correctly:

r = np.arange(A.shape[0]).reshape(-1, 1)
c = np.argsort(A)
A = A[r, c]
B = B[r, c]

The reason that your original index didn't work out is that you were indexing with a single dimension, which selects entire rows based on each location. This would have failed if you had more columns than rows.

A simpler way would be to follow what the argsort docs suggest:

A = np.take_along_axis(A, mask, axis=-1)
B = np.take_along_axis(B, mask, axis=-1)