Как отсортировать матрицу numpy с помощью маски?

Question

У меня есть две матрицы A, B, которые выглядят так:

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]])

Я пытаюсь отсортировать A и B НО Мне нужно, чтобы B был заказан с маской из A.

Я пробовал это:

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

Однако возвращаемое значение представляет собой матрицу (5, 5, 5)

Следующий фрагмент работает, но использует две итерации. Мне нужно что-то быстрее. У кого-нибудь есть идея?

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

Answer 1

Вы можете использовать необычную индексацию. Результат будет той же формы, что и ваши индексы, транслируемые вместе. Ваш индекс столбца уже правильной формы. Индекс строки размером (A.shape[0], 1) будет транслироваться правильно:

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

Причина, по которой исходный индекс не сработал, заключается в том, что вы индексировали с одним измерением, которое выбирает целые строки на основе каждого местоположения , Это не помогло бы, если бы у вас было больше столбцов, чем строк.

Проще было бы следовать указаниям документов argsort:

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