CODEWITHSUNDEEP
pythonarraysnumpysorting
user90465
user90465

Reputation: 187

How do I sort a 2D numpy array in this specific way

I realize there are quite a number of 'how to sort numpy array'-questions on here already. But I could not find how to do it in this specific way.

I have an array similar to this:

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

I want to sort the rows, keeping the order within the rows the same. So I expect the following output:

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

Upvotes: 2

Views: 85

Answers (1)

cs95
cs95

Reputation: 402253

You can use dot and argsort:

a[a.dot(2**np.arange(a.shape[1])[::-1]).argsort()]
# array([[0, 0, 1],
#        [1, 0, 1],
#        [1, 1, 0],
#        [1, 1, 1]])

The idea is to convert the rows into integers.

a.dot(2**np.arange(a.shape[1])[::-1])
# array([5, 1, 7, 6])

Then, find the sorted indices and use that to reorder a:

a.dot(2**np.arange(a.shape[1])[::-1]).argsort()
# array([1, 0, 3, 2])

My tests show this is slightly faster than lexsort.

a = a.repeat(1000, axis=0)

%timeit a[np.lexsort(a.T[::-1])]
%timeit a[a.dot(2**np.arange(a.shape[1])[::-1]).argsort()]

230 µs ± 18.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
192 µs ± 4.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Verify correctness:

np.array_equal(a[a.dot(2**np.arange(a.shape[1])[::-1]).argsort()], 
               a[np.lexsort(a.T[::-1])])
# True

Upvotes: 3

Related Questions