Как эффективно отранжировать множество векторов по близости к вектору-запросу?

«Как эффективно отранжировать множество векторов по близости к вектору-запросу?» — вопрос из категории Алгоритмы и структуры данных, который задают на 26% собеседований Data Scientist / ML Инженер. Ниже — развёрнутый ответ с разбором ключевых моментов.

Ответ

Выбор метода зависит от размерности векторов, их количества и требований к точности/скорости.

1. Точный поиск (Exact Search) для небольших коллекций Если векторов до сотен тысяч, я вычисляю все попарные расстояния и сортирую. Для этого использую оптимизированные операции с numpy.

import numpy as np

# Допустим, у нас есть матрица векторов и один запрос
vectors = np.random.rand(10000, 128)  # 10к векторов размерности 128
query = np.random.rand(128)

# Нормализуем для косинусного сходства (опционально, но часто нужно)
vectors_norm = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
query_norm = query / np.linalg.norm(query)

# Вычисляем косинусное сходство (1 = идентичные, -1 = противоположные)
cosine_similarities = np.dot(vectors_norm, query_norm)

# Получаем индексы векторов, отсортированных по убыванию сходства
most_similar_indices = np.argsort(cosine_similarities)[::-1]
top_5 = most_similar_indices[:5]

2. Приближенный поиск ближайшего соседа (ANN) для больших коллекций Для миллионов векторов точный поиск становится непрактичным. В этом случае я применяю ANN-библиотеки, которые жертвуют минимальной точностью ради огромного прироста скорости.

  • FAISS (Facebook AI Similarity Search): Использую для плотных векторов, особенно эффективен с GPU. Поддерживает индексы типа IVFFlat (кластеризация) или HNSW (графовый метод).
  • Annoy (Approximate Nearest Neighbors Oh Yeah): От Spotify, простая в использовании, строит леса случайных проекционных деревьев.
import faiss

# Предположим, vectors — это наша большая матрица векторов
d = vectors.shape[1]  # размерность
index = faiss.IndexFlatL2(d)  # Индекс для точного поиска по L2-расстоянию
# Или для приближенного:
# nlist = 100  # количество кластеров
# quantizer = faiss.IndexFlatL2(d)
# index = faiss.IndexIVFFlat(quantizer, d, nlist)
# index.train(vectors)

index.add(vectors)  # добавляем базу векторов в индекс

# Поиск k ближайших соседей
k = 10
distances, indices = index.search(query.reshape(1, -1), k)

Мой выбор: для точности в офлайн-задачах — точный поиск; для онлайн-сервисов с большими данными — FAISS с HNSW или IVFPQ.