Векторизованный поиск индексов в NumPy без O(n^2)

Как в NumPy найти индексы, где есть партнёр с f(x1,x2)>0, без декартова произведения: сводим условие к порогам по экстремумам и используем векторизацию.

Когда нужно просканировать два больших массива NumPy и отметить каждый индекс, у которого найдётся хотя бы один партнёр, удовлетворяющий условию, наивное декартово произведение — не вариант. Рассмотрим задачу: «для каждого элемента одного массива существует ли элемент в другом массиве такой, что f(x1, x2) > 0?» Реализация с циклом по одному массиву и векторизацией по другому работает, но плохо масштабируется и подталкивает к созданию матрицы 10 000 × 10 000, которая на самом деле не нужна.

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

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

import numpy as np

arr_left = np.random.rand(10000)
arr_right = np.random.rand(10000)

def score(u, v):
    return np.exp(v - u) / 1.2 - 1

ok_left = set()
ok_right = set()
for idx_l in range(len(arr_left)):
    idx_r_hits = np.nonzero(score(arr_left[idx_l], arr_right) > 0)[0].tolist()
    ok_right.update(idx_r_hits)
    if len(idx_r_hits) > 0:
        ok_left.add(idx_l)

print(sorted(ok_left))
print(sorted(ok_right))

Почему это медленно и что действительно важно

Ключ — в самом условии. Для функции f, определённой как exp(x2 − x1)/1.2 − 1, неравенство f(x1, x2) > 0 эквивалентно x2 > x1 + log(1.2). Значит, проверять каждую пару не нужно. Чтобы значение из первого массива было «валидным», достаточно, чтобы оно было достаточно маленьким: тогда найдётся хотя бы одно значение из второго массива, которое превысит его более чем на log(1.2). И наоборот, значение из второго массива «валидно», если оно достаточно велико, чтобы превзойти хотя бы одно значение из первого массива больше чем на тот же сдвиг.

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

Векторизованное решение

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

import numpy as np

arr_left = np.random.rand(10000)
arr_right = np.random.rand(10000)

delta = np.log(1.2)

idx_left = np.flatnonzero(arr_left < arr_right.max() - delta)
idx_right = np.flatnonzero(arr_right > arr_left.min() + delta)

Преобразование f(x1, x2) > 0 ⇔ x2 > x1 + log(1.2) даёт два простых порога. Элемент левого массива подходит, если он меньше max(arr_right) − log(1.2). Элемент правого массива подходит, если он больше min(arr_left) + log(1.2). Функция np.flatnonzero извлекает индексы за один проход по каждому массиву.

Как проверить корректность

Не пытайтесь проверять, попарно сопоставляя отфильтрованные индексы «один к одному» и заново оценивая f на этих парах. Исходная логика — экзистенциальная: «для элемента одного массива существует ли элемент в другом массиве, удовлетворяющий условию?» Попарная проверка игнорирует эту экзистенциальность и может ввести в заблуждение.

Надёжный способ — сравнить итоговые индексы, полученные быстрым методом, с индексами из исходной версии на циклах. Если множества совпадают, оптимизация корректна. Это можно утвердить с помощью проверок равенства NumPy:

# Вычислите ok_left / ok_right с помощью приведённой выше версии на циклах

# Вычислите idx_left / idx_right по векторизованным порогам из примера выше

left_match = np.array_equal(np.sort(np.fromiter(ok_left, dtype=int)), np.sort(idx_left))
right_match = np.array_equal(np.sort(np.fromiter(ok_right, dtype=int)), np.sort(idx_right))
print(left_match, right_match)

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

Такой подход избавляет от построения матрицы 10 000 × 10 000 и пропускает циклы на уровне Python по 10 000 элементам. Он использует алгебраическую форму условия, сводя задачу к простым сравнениям с предвычисленными экстремумами — именно в этом NumPy силён. Есть и более общие выводы. Для монотонных функций вроде exp сортировка и двоичный поиск часто сокращают объём работы, когда нужно найти минимальные удовлетворяющие элементы, а стратегия разбиения на тайлы (tiling) помогает снизить трафик памяти при матричных сканированиях. Важно и то, что не существует универсального аппаратного ускорения, которое «магически» разгонит проверку всех попарных комбинаций; реальный прирост даёт структурное понимание, как в неравенстве выше.

Вывод

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