Сортировка 2D numpy выстраивает несколькими осями
У меня есть 2D numpy массив формы (N, 2), который содержит точки N (координаты X и Y). Например:
array([[3, 2],
[6, 2],
[3, 6],
[3, 4],
[5, 3]])
Я хотел бы отсортировать его таким образом, что мои точки заказаны x-координатой, и затем y в случаях, где координата x является тем же. Таким образом, массив выше должен быть похожим на это:
array([[3, 2],
[3, 4],
[3, 6],
[5, 3],
[6, 2]])
Если бы это было нормальным списком Python, то я просто определил бы компаратор, чтобы сделать то, что я хочу, но насколько я могу сказать, функция вида numpy не принимает пользовательские компараторы. Какие-либо идеи?
Править: Спасибо за идеи! Я настроил быстрый тестовый сценарий с 1 000 000 случайных целочисленных точек и сравнил тех, что я мог работать (извините, не может обновить numpy в данный момент).
Mine: 4.078 secs
mtrw: 7.046 secs
unutbu: 0.453 secs
4 ответа
Использование lexsort :
import numpy as np
a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])
ind = np.lexsort((a[:,1],a[:,0]))
a[ind]
# array([[3, 2],
# [3, 4],
# [3, 6],
# [5, 3],
# [6, 2]])
a.ravel () возвращает представление, если a равно ПОСТОЯННЫЙ . Если это правда, то
метод @ ars , слегка измененный с помощью ravel вместо flatten , дает удобный способ сортировки a на месте :
a = np.array([(3, 2), (6, 2), (3, 6), (3, 4), (5, 3)])
dt = [('col1', a.dtype),('col2', a.dtype)]
assert a.flags['C_CONTIGUOUS']
b = a.ravel().view(dt)
b.sort(order=['col1','col2'])
Поскольку b является представлением a , сортировка b сортирует a а также:
print(a)
# [[3 2]
# [3 4]
# [3 6]
# [5 3]
# [6 2]]
Вы можете использовать np.complex_sort . Это имеет побочный эффект преобразования ваших данных в числа с плавающей запятой, надеюсь, это не проблема:
>>> a = np.array([[3, 2], [6, 2], [3, 6], [3, 4], [5, 3]])
>>> atmp = np.sort_complex(a[:,0] + a[:,1]*1j)
>>> b = np.array([[np.real(x), np.imag(x)] for x in atmp])
>>> b
array([[ 3., 2.],
[ 3., 4.],
[ 3., 6.],
[ 5., 3.],
[ 6., 2.]])
РЕДАКТИРОВАТЬ: удален плохой ответ.
Вот один из способов сделать это с помощью промежуточного структурированного массива:
from numpy import array
a = array([[3, 2], [6, 2], [3, 6], [3, 4], [5, 3]])
b = a.flatten()
b.dtype = [('x', '<i4'), ('y', '<i4')]
b.sort()
b.dtype = '<i4'
b.shape = a.shape
print b
, который дает желаемый результат:
[[3 2]
[3 4]
[3 6]
[5 3]
[6 2]]
Не уверен, что это лучший способ сделать это.
Я нашел один способ сделать это:
from numpy import array
a = array([(3,2),(6,2),(3,6),(3,4),(5,3)])
array(sorted(sorted(a,key=lambda e:e[1]),key=lambda e:e[0]))
Довольно ужасно приходится сортировать дважды (и использовать простой python sorted вместо более быстрой сортировки numpy), но она прекрасно помещается в одной строке.