Сортировка 2D массива по первым n строкам

Question

Как мне отсортировать массив в NumPy по двум первым строкам?

Например,

A=array([[9, 2, 2],
         [4, 5, 6],
         [7, 0, 5]])

И я бы хотел отсортировать столбцы по первым двум строкам такой, что я вернусь:

A=array([[2, 2, 9],
         [5, 6, 4],
         [0, 5, 7]])

Спасибо!

Answer 1

Один из подходов заключается в преобразовании двумерного массива, над которым мы хотим взять argsort, в более простой в обращении одномерный массив. Для этого одной идеей может быть умножение строк для учета в целях сортировки путем последовательного уменьшения значений в степени 10 последовательности, sum их и , затем использования argsort ( примечание : этот метод будет численно нестабильным для больших значений k . Предназначен для значений до ~ 20 ):

def sort_on_first_k_rows(x, k):
    # normalize each row so that its max value is 1
    a = (x[:k,:]/x[:k,:,None].max(1)).astype('float64')
    # multiply each row by the seq 10^n, for n=k-1,k-2...0
    # Ensures that the contribution of each row in the sorting is
    # captured in the final sum
    a_pow = (a*10**np.arange(a.shape[0]-1,-1,-1)[:,None])
    # Sort with the argsort on the resulting sum
    return x[:,a_pow.sum(0).argsort()]

Проверка с помощью общий пример:

sort_on_first_k_rows(A, 2)
array([[2, 2, 9],
       [5, 6, 4],
       [0, 5, 7]])

---

Или с другим примером:

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

sort_on_first_k_rows(A, 2)
array([[1, 2, 2, 2, 5, 9, 9],
       [0, 3, 6, 7, 9, 3, 4],
       [0, 1, 5, 0, 2, 2, 7]])

Answer 2

Библиотека pandas очень гибкая для сортировки DataFrames - но только на основе столбцов. Поэтому я предлагаю транспонировать и преобразовывать ваш массив в DataFrame следующим образом (обратите внимание, что вам необходимо указать имена столбцов для последующего определения критериев сортировки):

df = pd.DataFrame(A.transpose(), columns=['col'+str(i) for i in range(len(A))])

Затем сортируйте и конвертируйте обратно, как это:

A_new = df.sort_values(['col0', 'col1'], ascending=[True, True]).to_numpy().transpose()