有效地找到大图的近似最小生成树

我有大量节点，大约 25000 个，每个节点都有 3D 位置。

我想生成一个稀疏连接的图，其边由节点之间的距离给出，以在 GNN 中使用。

查找算法最小生成树 (MST) 通常依赖于首先从完全连接的图开始，然后删除节点以找到 MST。对于大型图来说，这非常慢。

为了尝试加快速度，我尝试使用 scipy 的稀疏距离矩阵将初始连接的半径限制为最大最近邻距离，但这会导致某些图出现多个连接组件。|| |（此方法的较小图的示例不起作用：）

这是我尝试过的：

我不需要真正的最小生成树，只是为了使图与尽可能少的连接边缘以加速 GNN 性能。即使稀疏方法对于某些图来说也太慢了。

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)

我考虑了一种基于

https://www.cs.umd.edu/~mount/Papers/soda16-emst.pdf 的方法，但它看起来很难实现，即 scipy 没有四叉树。 将全距离矩阵转换为 networkx 图，然后使用他们的 Boruvka 算法的实现甚至更慢，它不适用于大型图。 向 max_neighbor_distance 添加一个乘数将有助于确保只有一个连接的组件，但也会增加处理时间，并且并不总是足够的。

Converting the fully distance matrix to a networkx graph then using their implementation of Boruvka's algorithm is even slower, it's not intended for large graphs. Adding a multiplier to max_neighbor_distance would help ensure there's only one connected component, but would also increase processing time and won't always be enough.

你遇到的问题很常见：你想要一个大型数据集的最小生成树的快速近似，以用于你的 GNN。虽然完整的 MST 计算成本高昂，但你是对的，仅基于最大近邻距离进行阈值化可能会导致不相交的组件。

让我们探讨一些可以帮助你实现速度和连通性的策略：

1. 混合方法：K 近邻图 + MST

2. 理念： 与其依赖单个阈值，不如构建一个保证连通性的初始图，然后对其进行修剪以获得稀疏性。

3. 步骤：
   1. K 近邻图： 为你的数据集构建一个 K 近邻图。你可以使用 KDTree.query(k=K) 来有效地找到每个点的 K 个最近邻。选择一个小的 K 值（例如，3 到 5）以确保连通性，但不要太大以保持稀疏性。
   2. MST 计算： 在 K 近邻图上计算 MST。由于这个图已经比完全图稀疏得多，因此 MST 计算应该快得多。
4. 优点： 保证连通性，并且比完全 MST 快。

5. 缺点： MST 可能不是真正全局最优的，因为你从一个更密集的初始图开始。

6. 增量边添加（类似于 Prim 算法）

7. 理念： 以一个小的连通图开始，并迭代地添加连接现有图和剩余节点的最小成本边。

8. 步骤：
   1. 初始化： 从一个包含单个节点的图开始。
   2. 迭代增长： 在连接现有图中的节点和图外的节点的边中找到成本最低的边。将这条边和相应的节点添加到图中。
   3. 终止： 重复步骤 2，直到所有节点都已连接。
9. 优点： 逐渐构建 MST，避免了大型密集图。

10. 缺点： 可能不如 K 近邻 + MST 方法快，具体取决于实现。

11. 基于四叉树的近似（如果你需要非常大的规模）

12. 理念： 虽然 scipy 没有直接的四叉树实现，但你可以使用其他库或探索近似方法。基于四叉树的分解可以帮助你递归地划分你的空间并找到连接不同区域的近似最小成本边。

13. 库： 请查看 scikit-learn 中的 BallTree ，它可以很好地替代四叉树，并且在某些情况下性能甚至更好。
14. 复杂性： 此方法的实现更加复杂，但对于非常大的数据集和需要严格控制近似误差的情况非常有用。

代码示例（K 近邻图 + MST）：

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

# 创建一组随机的 25000 个节点位置
positions = np.random.normal(loc=500000, scale=10000, size=(25000, 3))

# 步骤 1：K 近邻图
k = 5  # 选择一个小 k 值
kdtree = KDTree(positions)
distances, indices = kdtree.query(positions, k=k + 1)  # 查询 k+1 个邻居，包括自身

# 从索引和距离创建稀疏矩阵
row_indices = np.repeat(np.arange(len(positions)), k)
col_indices = indices[:, 1:].flatten()  # 排除自身
data = distances[:, 1:].flatten()
sparse_dist_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(len(positions), len(positions)))

# 步骤 2：MST 计算
mst = minimum_spanning_tree(sparse_dist_matrix)

# 创建 networkx 图以供使用
G = nx.from_scipy_sparse_array(mst)

提示：

• 基准测试： 尝试不同的方法，并在你的特定数据集上对它们进行基准测试，以查看哪种方法提供最佳性能。
• 参数调整： 诸如 K 值（对于 K 近邻）或四叉树分解中的误差容限等参数会影响结果的稀疏性和准确性。

请记住，找到完美的 MST 近似是一个权衡问题。目标是找到一种方法，在速度和保持图的连通性之间取得良好的平衡，以满足你的 GNN 应用程序的需求。