Эффективно найдите приблизительное минимальное остовное дерево большого графа

У меня большое количество узлов, ~25000, каждый из которых имеет трехмерную позицию.

Я хочу создать разреженный граф с ребрами, определяемыми расстоянием между узлами, для использования в GNN.

Алгоритмы поиска минимального связующего дерева (MST) обычно основаны на том, что сначала начинается с полностью связного графа, а затем удаляются узлы для поиска MST. Это очень медленно для больших графиков.

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

(Пример меньшего графика этого метода не работает:)

Вот что я пробовал:

import numpy as np
from scipy.sparse.csgraph import minimum_spanning_tree
from scipy.spatial import KDTree
from scipy.spatial import distance_matrix

# Create a random set of 25000 node positions
positions = np.random.normal(loc=500000, scale=10000, size=(25000, 3))

# Method 1: Full MST Search (Too Slow)
dist_matrix = distance_matrix(positions, positions)
mst = minimum_spanning_tree(dist_matrix)


# Method 2: Start from sparse matrix (Results in multiple connected components)
kdtree = KDTree(positions)

# Find the maximum nearest neighbour distance for the graph
distances, _ = kdtree.query(positions, k=2)
max_neighbour_distance = np.max(distances[:, 1])
max_neighbour_distance = np.ceil(max_neighbour_distance) # Round up to avoid floating point errors in MST search


# Threshold the distance matrix by this distance
sparse_dist_matrix = kdtree.sparse_distance_matrix(kdtree, max_neighbour_distance, output_type = "coo_matrix")

mst = minimum_spanning_tree(sparse_dist_matrix)

G = nx.from_scipy_sparse_array(mst)

Мне не нужно истинное минимальное остовное дерево, просто чтобы граф был связан с как можно меньшим количеством ребер, чтобы ускорить производительность GNN. Даже разреженный метод слишком медленный для некоторых графов.

Я рассматривал метод, основанный на https://www.cs.umd.edu/~mount/Papers/soda16-emst.pdf, но его сложно реализовать, т. е. в scipy нет квадродеревьев.

Преобразование полной матрицы расстояний в граф networkx с последующим использованием реализации алгоритма Борувки происходит еще медленнее, оно не предназначено для больших графов. Добавление множителя к max_neighbor_distance поможет гарантировать наличие только одного подключенного компонента, но также увеличит время обработки, которого не всегда будет достаточно.