python - 使用 K-D 树的 NetworkX 随机几何图实现

Question

所以很明显，NetworkX 是用一个算法在 n^2 次生成一个随机的几何图形。他们说使用 K-D 树可以实现更快的算法。我的问题是如何尝试实现该算法的 K-D 树版本？我不熟悉这种数据结构，也不会称自己为 python 专家。只是想弄清楚这一点。感谢所有帮助，谢谢!

    def random_geometric_graph(n, radius, dim=2, pos=None):
        G=nx.Graph()
        G.name="Random Geometric Graph"
        G.add_nodes_from(range(n)) 
        if pos is None:
            # random positions
            for n in G:
                G.node[n]['pos']=[random.random() for i in range(0,dim)]
        else:
            nx.set_node_attributes(G,'pos',pos)
        # connect nodes within "radius" of each other
        # n^2 algorithm, could use a k-d tree implementation
        nodes = G.nodes(data=True)
        while nodes:
            u,du = nodes.pop()
            pu = du['pos']
            for v,dv in nodes:
                pv = dv['pos']
                d = sum(((a-b)**2 for a,b in zip(pu,pv)))
                if d <= radius**2:
                    G.add_edge(u,v)
    return G

Answer 1

这是一种使用上面@tcaswell 提到的 scipy KD-tree 实现的方法。

import numpy as np
from scipy import spatial
import networkx as nx
import matplotlib.pyplot as plt
nnodes = 100
r = 0.15
positions =  np.random.rand(nnodes,2)
kdtree = spatial.KDTree(positions)
pairs = kdtree.query_pairs(r)
G = nx.Graph()
G.add_nodes_from(range(nnodes))
G.add_edges_from(list(pairs))
pos = dict(zip(range(nnodes),positions))
nx.draw(G,pos)
plt.show()