Почему argsort(argsort(x)) возвращает ранги элементов массива / Хабр

Однажды мне понадобилось вычислить ранги элементов массива, или другими словами, на какой позиции окажется каждый элемент после сортировки.

Беглый поиск по интернету дал неожиданно короткий ответ: достаточно применить двойной argsort —argsort(argsort(x)) из библиотеки NumPy. Фокус работает, но почти нигде не объясняется почему именно. Обращаться с таким вопросом к ChatGPT тогда казалось рискованным — в интернете как раз ходили мемы о его ошибках в арифметике. Отсутствие строгого обоснования не давало мне покоя, и в итоге я решил дойти до доказательства самостоятельно.

Ниже я поделюсь этим доказательством. Оно не претендует на уникальность — возможно, существуют более короткие или более элегантные способы прийти к тому же ответу. Тем не менее, оно позволит вам понять, почему двойной argsort действительно возвращает ранги элементов, и избавит вас от необходимости тратить время на самостоятельное разбирательство.

На самом деле в Python в библиотеке scipy уже есть функция rankdata, которая реализует то, что нам нужно. Единственное отличие — rankdata возвращает ранги начиная с 1, тогда как argsort использует нулевую индексацию. Чуть ниже мы сравним их результаты, а пока формализуем нашу задачу.

Основные определения

Пусть задан массив:

$a = (a_1, a_2, \dots, a_N)$

$s = (s_1, s_2, \dots, s_N)$ - его отсортированная версия

$m = (m_1, m_2, \dots, m_N)$ - результат работы argsort(a)

$p = (p_1, p_2, \dots, p_N)$ - результат работы argsort(argsort(a))

$r = (r_1, r_2, \dots, r_N)$ - ранги массива

Отдельно отмечу, что если $a_i \neq a_k$ для $\forall i \neq k$ , то также показывает, сколько элементов массива меньше данного, ведь для $\forall i\geq1$ верно следующее $s_{r_{i-1}} < s_{r_i}$ . Если же в массиве присутствуют одинаковые элементы, то встает вопрос, какой ранг им присвоить. Двойной аргсорт будет выдавать ранг по исходной позиции элемента: чем меньше индекс в исходном массиве, тем меньше итоговый ранг. Поэтому значение ранга у него уникально для каждого элемента, но оно уже не равно количеству элементов, меньших данного. В rankdata из scipy можно настроить этот момент под конкретную задачу через переменную method.

Для наглядности рассмотрим пример

Чтобы получить от rankdata такой же ответ как от argsort мы укажем в качестве method='ordinal' а также вычтем единицу, чтобы его ответы лежали в том же множестве.

import numpy as np
from scipy.stats import rankdata


a = np.array([13, 0, 47, 52, 27])

print("a:", a)

s = np.sort(a)
print("s = sort(a):", s)

m = np.argsort(a)
print("m = argsort(a):", m)

p = np.argsort(m)
print("p = argsort(argsort(a)):", p)

r = rankdata(a, method="ordinal") - 1
print("r = rank(a):", r)

a: [13 0 47 52 27]
s = sort(a): [0 13 27 47 52]
m = argsort(a): [1 0 4 2 3]
p = argsort(argsort(a)): [1 0 3 4 2]
r = rank(a): [1 0 3 4 2]

Из определения argsort следует, что $a_{m_i} = s_i, \forall i \in \{1,\dots,N\}$ , т.е. - это индекс элемента массива a, который после сортировки будет на позиции i в отсортированном массиве s.

Мы же хотим найти ранг , который показывает позицию (индекс), которую элемент занимает после сортировки. В терминах введенных выше обозначений $s_{r_i} = a_i, \forall i \in \{1,\dots,N\}$ .

Утверждается, что $p_i = r_i, \forall i \in \{1,\dots,N\}$ , докажем это.

Доказательство

Из определения argsort следует, что массив m является перестановкой индексов $\{1,\dots,N\}$ .

То есть:

$m_i \in \{1,\dots,N\}, \quad m_i \ne m_k \ \text{для } i \ne k \quad(1)$

Рассмотрим массив:

По определению argsort индексы p упорядочивают элементы массива m по не убыванию:

$m_{p_1} \le m_{p_2} \le \dots \le m_{p_N}$

Но из условия (1) мы знаем, что равенство невозможно, и итоговая последовательность будет представлять из себя упорядоченные индексы от 1 до , т.е. $m_{p_1} = 1, m_{p_2} = 2, \dots, m_{p_N} = N$