Эффективная группировка в NumPy: unique(return_inverse) + add.at вместо bincount

Как сгруппировать пары по ключу в NumPy без гигантских массивов и поисков: unique(return_inverse) + numpy.add.at вместо bincount. Для разрежённых ключей.

Когда у вас есть ndarray из пар вида [curveLocations, distinctCrossings], и некоторые curveLocations повторяются, задача — сгруппировать строки по этому ключу и просуммировать связанные счётчики. Сортировка не обязательна, а вот быстрое, векторизованное группирование — да. Часто пытаются использовать numpy.bincount: выглядит заманчиво, но для разреженных ключей с большими значениями это может обернуться бедой. В этой задаче лучше подходит связка numpy.unique с параметром return_inverse и последующим «рассеянным» сложением через numpy.add.at.

Постановка задачи

На вход подаётся неотсортированный двумерный ndarray из беззнаковых целых. В первом столбце — curveLocations (ключи), во втором — distinctCrossings (значения для суммирования). Если ключ повторяется, его значения нужно агрегировать. Сортировать результат не обязательно.

Ключи разрежённые, а максимальное значение ключа может быть очень большим. Как сказано в документации,

«Длина out равна np.amax(x)+1.»

Из‑за этого numpy.bincount здесь плохо подходит: с ростом максимума число корзин взрывообразно увеличивается, при том что подавляющая их часть пустует. В рассматриваемом сценарии максимум ключа ограничен curveLocationsMAXIMUM = 1 << (2 * n + 4). При n = 25 получаем почти 2^54 корзин — очевидно, это неприемлемо.

Пример кода, иллюстрирующий проблему

Ниже — векторизованное объединение с использованием отсортированного набора уникальных ключей и отдельного шага поиска для каждой входной строки. Это отражает подход, в котором numpy.unique используется для извлечения ключей, а numpy.searchsorted — для сопоставления каждой исходной строки с индексом её группы. Имена переменных выбраны для наглядности; поведение соответствует известному шаблону.

import numpy as np
def pack_by_location(pairs2d: np.ndarray) -> np.ndarray:
    # pairs2d имеет форму (N, 2): [:, 0] — ключи (curveLocations), [:, 1] — значения (distinctCrossings)
    unique_keys = np.unique(pairs2d[:, 0])
    # Готовим результат из 2 столбцов: [ключ, сумма]
    grouped = np.tile(unique_keys, (2, 1)).T
    grouped[:, 1] = 0
    # Отобразим каждый входной ключ на его позицию в unique_keys, затем выполним «рассеянное» сложение значений
    idx_map = np.searchsorted(grouped[:, 0], pairs2d[:, 0])
    np.add.at(grouped[:, 1], idx_map, pairs2d[:, 1])
    return grouped

Здесь выполняются два векторизованных прохода: первый — чтобы извлечь и отсортировать ключи, второй — чтобы найти индекс каждой исходной строки среди этих ключей. Это работает, но связка numpy.unique и numpy.searchsorted может «съедать» основное время выполнения — как показало профилирование, большая часть затрат в aggregateCurveLocations приходилась именно на эти функции.

Почему так происходит

Само объединение предельно простое: нужно сопоставить каждой исходной строке индекс группы и добавить её значение в соответствующее «ведро». Важный нюанс — как эффективно построить это отображение. numpy.bincount неприменим, потому что выделяет массив до максимального ключа — а он колоссален при разрежённых и широких пространствах ключей. Поисковый подход из предыдущего примера корректен, но требует дополнительного прохода бинарных поисков после получения отсортированных уникальных ключей.

numpy.unique умеет сразу дать нужное отображение через параметр return_inverse. Этот «обратный» массив — прямой индекс в массив уникальных ключей для каждого элемента входа, так что отдельный этап поиска не нужен. Получив индексы, используйте numpy.add.at для эффективного «рассеянного» сложения и накопления сумм.

Решение

Функция ниже объединяет по ключу с помощью numpy.unique(return_inverse=True) и numpy.add.at. Она избавляется от лишнего прохода с numpy.searchsorted и сохраняет компактное потребление памяти, не зависящее от величины ключей.

import numpy as np
def collapse_pairs(pairs2d: np.ndarray) -> np.ndarray:
    # pairs2d имеет форму (N, 2): [:, 0] — ключи (curveLocations), [:, 1] — значения (distinctCrossings)
    keys, rev_idx = np.unique(pairs2d[:, 0], return_inverse=True)
    totals = np.zeros_like(keys)
    np.add.at(totals, rev_idx, pairs2d[:, 1])
    return np.column_stack((keys, totals))

На выходе получается двумерный массив из двух столбцов: [curveLocations, summedDistinctCrossings]. По умолчанию numpy.unique возвращает ключи в отсортированном виде — это подходит, поскольку сохранять исходный порядок здесь не требуется.

Расширение за рамки двух столбцов

Если после объединения нужны дополнительные столбцы, вычисляйте их по сгруппированному результату. Например, побитовые извлечения вроде groupZulu и groupAlpha можно получить из консолидированных curveLocations теми же масками и сдвигами, что и раньше. Шаг объединения ортогонален этим вычислениям и лишь заменяет более медленную «группировку через поиск».

Почему это стоит держать под рукой

Группировка по ключу с суммированием — базовая операция в задачах с интенсивной обработкой данных, и важно уметь делать её правильно на уровне ndarray. Разрежённые пространства ключей исключают трюки с подсчётом по корзинам; построчные циклы Python не вариант; повторные поиски создают лишние накладные расходы. Связка unique + inverse + add.at компактна, не требует гигантских выделений памяти и стабильно масштабируется на больших неотсортированных входах.

Выводы

Для «удаления дублей через объединение» в NumPy формулируйте задачу так: «построить для каждой строки индекс группы и выполнить рассеянное сложение». Когда ключи разрежённые и потенциально огромные, откажитесь от numpy.bincount. Используйте numpy.unique с return_inverse=True, чтобы сразу получить индексы групп, и накапливайте суммы через numpy.add.at. Если конвейер включает постагрегационные поля, вычисляемые из curveLocations, рассчитайте их уже по сгруппированному результату — точно так же, как раньше.