Numpy sort N dimensional array similar to pythons list sort

Question

I am trying to sort a list of 3D coordinates. The target is to get the same order that the basic Python list sort produces, which is the same as the meshgrid + column_stack order. Python sorts by the first value of each element first; in case of ties, the second element is checked, and so on. I would like to do the same in numpy, especially with the argsort function, as it gives an index list for the sorted order. I need this list to sort other numpy arrays with the same first dimension.

I think this post is related, but does not use numpy to do the sorting.

Following a small test bench.

import numpy as np

num = 10
x = np.linspace(1, 3, num)
y = np.linspace(4, 6, num)
z = np.linspace(7, 9, num)
xg, yg , zg = np.meshgrid(x, y, z, indexing='ij', sparse=False)
test = np.column_stack((xg.ravel(), yg.ravel(), zg.ravel()))
print('Target ordering', test)

np.random.shuffle(test)
print('\nShuffeld', test)

python_sorted = sorted(test.tolist())
python_sorted = np.array(python_sorted)
print('\nPython list sorted', python_sorted)

numpy_sort_idx = test.argsort(axis=0)
numpy_sorted = test[numpy_sort_idx]
print('\nNumpy arg sorted', numpy_sorted)

Simon T. · Accepted Answer · 2025-04-18 10:04:13Z

2

You could use lexsort on the reversed columns:

import numpy as np

num = 10
x = np.linspace(1, 3, num)
y = np.linspace(4, 6, num)
z = np.linspace(7, 9, num)
xg, yg , zg = np.meshgrid(x, y, z, indexing='ij', sparse=False)
test = np.column_stack((xg.ravel(), yg.ravel(), zg.ravel()))
print('Target ordering\n', test)

np.random.shuffle(test)
print('\nShuffled\n', test)

python_sorted = sorted(test.tolist())
python_sorted = np.array(python_sorted)
print('\nPython list sorted\n', python_sorted)

idx = np.lexsort(test[:, ::-1].T)
numpy_sorted = test[idx]
print('\nLexsorted using reversed columns\n', numpy_sorted)

The order of the columns is reversed due to the precedence of the keys passed to lexsort. As the first column should be the primary sort order, it must come last. From the documentation:

The last key in the sequence is used for the primary sort order, ties are broken by the second-to-last key, and so on.

Output form the test bench

Target ordering
 [[1.         4.         7.        ]
 [1.         4.         7.22222222]
 [1.         4.         7.44444444]
 ...
 [3.         6.         8.55555556]
 [3.         6.         8.77777778]
 [3.         6.         9.        ]]

Shuffled
 [[3.         5.11111111 7.22222222]
 [2.77777778 5.11111111 8.11111111]
 [1.66666667 5.77777778 7.22222222]
 ...
 [2.55555556 4.88888889 7.22222222]
 [2.55555556 4.88888889 7.66666667]
 [2.55555556 4.66666667 7.66666667]]

Python list sorted
 [[1.         4.         7.        ]
 [1.         4.         7.22222222]
 [1.         4.         7.44444444]
 ...
 [3.         6.         8.55555556]
 [3.         6.         8.77777778]
 [3.         6.         9.        ]]

Lexsorted using reversed columns
 [[1.         4.         7.        ]
 [1.         4.         7.22222222]
 [1.         4.         7.44444444]
 ...
 [3.         6.         8.55555556]
 [3.         6.         8.77777778]
 [3.         6.         9.        ]]

edited Apr 18 at 10:04

Simon T.

14511 bronze badges

answered Apr 16 at 12:31

Sash Sinha

22.8k4 gold badges27 silver badges48 bronze badges

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

3 Comments

Simon T. Apr 16 at 14:15

The lexsort looks good. I want to make sure. It says lexsort, but for numbers, it is doing a numerical sort by the actual float values? From the documentation, this does not seem obvious to me. As a lexicographic sort of numbers interpreter as strings can have unexpected results, e.g., when positive and negative numbers are mixed or 10 being "smaller" than 2. Hence, I want to clarify.

hpaulj Apr 16 at 19:04

Lexical sorting of your (n,3) array means treating each 'row' as a 'word'. Sort on the first 'letter/column', and break ties with the second, and so forth. Neither python nor numpy convert numbers to strings to compare them.

Simon T. Apr 17 at 16:11

Got it, makes sense. Thanks for clarifying.

Ahmed AEK · Accepted Answer · 2025-04-16 12:32:20Z

2

there's numpy.lexsort for this.

# sorts by test[:,0] then test[:,1] ...
numpy_sort_idx = np.lexsort((test[:,2], test[:,1], test[:,0]))
numpy_sorted = test[numpy_sort_idx]
print('\nNumpy arg sorted', numpy_sorted)

for arbitrary test.shape[1]

numpy_sort_idx = np.lexsort([test[:,i] for i in reversed(range(test.shape[1]))])

test on num = 2

Numpy arg sorted [[1. 4. 7.]
 [1. 4. 9.]
 [1. 6. 7.]
 [1. 6. 9.]
 [3. 4. 7.]
 [3. 4. 9.]
 [3. 6. 7.]
 [3. 6. 9.]]

edited Apr 16 at 12:32

answered Apr 16 at 12:26

Ahmed AEK

23.2k3 gold badges19 silver badges50 bronze badges

Collectives™ on Stack Overflow

Numpy sort N dimensional array similar to pythons list sort

2 Answers 2

Output form the test bench

3 Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Output form the test bench

3 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related