Быстрое сопоставление групп индексов в NumPy без object-массивов

Векторизуйте сопоставление групп индексов и маскированные обновления в NumPy без массивов object: np.repeat, данные и быстрые операции на больших массивах.

Наложение одного одномерного массива на другой через сопоставление групп целевых индексов кажется простой задачей — пока не приходится делать это на масштабе, с булевыми масками и частичными обновлениями. Типичная ловушка — хранить отображение во вложенных списках Python и затем в чистом Python жонглировать переводом индексов и расширением значений. При массивах порядка 100 тысяч элементов это быстро превращается в узкое место. Решение — отказаться от массивов с типом object и перенести весь конвейер в плоские, векторизованные операции NumPy.

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

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

import numpy as np

# управляющий массив
cover = np.array([0, 1, 1, 4, 3])

# отображение cover[i] -> позиции в целевом массиве
spans = [[0, 1, 2, 3], [4, 5], [6], [7], [8]]

# целевой массив (текущее состояние)
base = np.array([0, 0, 0, 0, 1, 1, 1, 4, 3])

# обновление: выбираем 1 и 3, присваиваем новые значения
pick = (cover == 1) | (cover == 3)
updates = np.array([44, 48, 47])
cover[pick] = updates

# наивный перевод индексов и расширение значений (неэффективно)
spans_obj = np.asarray(spans, dtype=object)
sel_spans = spans_obj[pick]
base_pos = sum(sel_spans, [])
expanded = np.repeat(updates, np.asarray([len(s) for s in sel_spans]))
base[base_pos] = expanded
print(base)  # ожидаемый результат: [0, 0, 0, 0, 44, 44, 48, 4, 47]

Это работает, но опирается на списки и тип object, что убивает векторизацию и добавляет накладные расходы, пропорциональные числу объектов Python, которыми вы манипулируете.

Почему это медленно

NumPy оптимизирован под непрерывную однородную память. Как только появляются рваные списки или массивы объектов, операции скатываются к циклам на уровне Python. Это и происходит, когда вы суммируете вложенные списки индексов и вручную повторяете значения по длине подсписков. Проблема не в алгоритме как таковом, а в представлении данных: вы кодируете плоское отображение через вложенные структуры Python вместо простого числового формата.

Векторизованный подход

Отображение от управляющего массива к целевому удобно представить через длины пробегов. Заранее посчитайте, сколько целевых индексов контролирует каждый элемент, и используйте np.repeat, чтобы и собрать целевой массив, и перевести булевы маски и значения за один проход. Так вы избегаете массивов объектов и сохраняете всё плоским и векторизованным.

import numpy as np

# управляющий массив
cover = np.array([0, 1, 1, 4, 3])

# отображение как группы индексов (каждый индекс base покрыт ровно один раз)
spans = [[0, 1, 2, 3], [4, 5], [6], [7], [8]]

# заранее считаем длины пробегов: сколько позиций контролирует каждый элемент
runs = np.array([len(g) for g in spans])

# собираем целевой массив один раз из cover по этому отображению
base = np.repeat(cover, runs)
print(base)  # [0 0 0 0 1 1 1 4 3]

# выбираем по булевой маске на cover
pick = (cover == 1) | (cover == 3)

# применяем обновления, транслируя маску и расширяя значения теми же длинами
base[np.repeat(pick, runs)] = np.repeat(np.array([44, 48, 47]), runs[pick])
print(base)  # [ 0  0  0  0 44 44 48  4 47]

Всё дорогое происходит один раз — при вычислении длин пробегов. Дальше мы остаёмся в мире векторизации: никаких циклов Python, никаких массивов объектов и конкатенаций «на лету».

Если удобнее мутировать управляющий массив и материализовать целевой только в конце — это так же просто. Полезно, когда над управляющим массивом идёт цепочка преобразований, а в целевое пространство нужно спроецировать результат один раз.

import numpy as np

cover = np.array([0, 1, 1, 4, 3])
spans = [[0, 1, 2, 3], [4, 5], [6], [7], [8]]
runs = np.array([len(g) for g in spans])

pick = (cover == 1) | (cover == 3)
cover[pick] = [44, 48, 47]

# материализуем один раз в конце
base = np.repeat(cover, runs)
print(base)  # [ 0  0  0  0 44 44 48  4 47]

Нужны ли вам вообще плоские индексы?

Когда группы покрытий охватывают каждый индекс в целевом массиве ровно один раз, «сплющенный» массив индексов всегда просто последовательность от нуля до последнего индекса. В этом случае использовать целочисленные позиции, полученные из такого массива, избыточно. Можно обращаться к нужным позициям напрямую через np.repeat для булевой маски, потому что повторённая маска имеет ту же длину и порядок, что и целевой массив.

Если всё же нужен явный плоский массив индексов, его можно построить один раз с помощью np.concatenate(spans). Но когда покрытие полное, такой массив совпадает с np.arange(len(base)). Примеры выше показывают, как безопасно обходиться без него.

Почему это важно

На массивах порядка 100 тысяч элементов массивы объектов и операции со списками создают ощутимые накладные расходы и лишнюю аллокацию. Представление отображения через длины пробегов сжимает уровень косвенности до числового вида, который NumPy умеет эффективно обрабатывать. В итоге задача сводится к паре операций repeat и одному маскированному присваиванию — всё остаётся дружелюбным к кешу и векторизованным.

Выводы

Держите отображение плоским. Заранее посчитайте, сколько целевых позиций контролирует каждый источник, и переиспользуйте эти длины. Транслируйте булевы маски через np.repeat и расширяйте значения с тем же срезом длин. Обновляете ли вы целевой массив на месте или пересобираете его из преобразованного управляющего — такой подход убирает Python с горячего пути и масштабируется предсказуемо.