Python实现不同半径连通球体分组高效检测（适配大规模数据）

问题说明 #

需要检测三维空间中不同半径球体的连通（接触/相交）关系，输出连通分组。给定测试数据如下：

import numpy as np
radii = np.array([2, 1, 1, 2, 2, 0.5])
poss = np.array([[7, 7, 7], [7.5, 8.5, 6], [0, 0, 0], [-1, -2, -1], [1, 1, 1], [2, 1, 3]])

期望输出支持两种格式：形状为(分组数量, 球体总数)的布尔数组，或是存储各分组对应球体索引的列表/数组。示例对应期望输出为：

Boolean_array = np.array([[1, 1, 0, 0, 0, 0], [0, 0, 1, 1, 1, 1]], dtype=bool)
object_array = np.array([[0, 1], [2, 3, 4, 5]])

原有实现问题 #

此前基于networkx编写的实现代码如下：

G = nx.Graph([*ends_ind])
L = [nx.node_connected_component(G, 0)]
for i in range(len(radii)):
    iter = 0
    for j in L:
        if i in j:
            iter += 1
    if iter == 0:
        L.append(nx.node_connected_component(G, i))

运行触发KeyError: 11报错，核心原因是传入构图的ends_ind包含超出球体索引范围的边：示例共6个球体，合法索引范围为0-5，但边数据中出现了不存在的节点索引11，导致遍历邻接表时查找失败。
除此之外，原有实现逐节点遍历所有已生成分量检查归属的逻辑时间复杂度极高，且networkx本身的图存储、遍历开销较大，完全无法支撑10万级球体的大规模计算场景。

10万级场景最优方案 #

整体实现分为两个核心步骤，全流程基于C优化的数值计算库实现，避免Python层面的循环开销：

• 碰撞对快速筛选：用KDTree做空间分区，跳过不可能碰撞的球对，避免O(n²)的全量两两距离计算。两球连通的判定规则为：球心欧氏距离 ≤ 两球半径之和。
• 连通分量合并：用并查集（Union-Find，不相交集合）结构做连通关系合并，操作近似O(1)时间复杂度，效率远高于networkx的图遍历方案。

完整实现代码 #

import numpy as np
from scipy.spatial import cKDTree

class UnionFind:
    def __init__(self, size):
        self.parent = np.arange(size)
        self.rank = np.zeros(size, dtype=np.int32)
    
    def find(self, x):
        # 路径压缩
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x
    
    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root == y_root:
            return
        # 按秩合并
        if self.rank[x_root] < self.rank[y_root]:
            self.parent[x_root] = y_root
        else:
            self.parent[y_root] = x_root
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1

def sphere_connectivity_group(poss, radii):
    n_sphere = len(radii)
    uf = UnionFind(n_sphere)
    # 构建球心KDTree索引
    kd_tree = cKDTree(poss)
    # 粗筛：所有球心距离小于2倍最大半径的球对才可能发生接触
    max_radius = radii.max()
    candidate_pairs = kd_tree.query_pairs(r=2*max_radius, output_type="ndarray")
    # 精筛：校验球对是否满足接触条件
    if len(candidate_pairs) > 0:
        pair_dist = np.linalg.norm(
            poss[candidate_pairs[:, 0]] - poss[candidate_pairs[:, 1]], 
            axis=1
        )
        radius_sum = radii[candidate_pairs[:, 0]] + radii[candidate_pairs[:, 1]]
        valid_pairs = candidate_pairs[pair_dist <= radius_sum]
        # 合并连通节点
        for i, j in valid_pairs:
            uf.union(i, j)
    # 整理分组结果
    root_to_group = dict()
    groups = []
    for idx in range(n_sphere):
        root = uf.find(idx)
        if root not in root_to_group:
            root_to_group[root] = len(groups)
            groups.append([])
        groups[root_to_group[root]].append(idx)
    # 生成布尔格式结果
    bool_result = np.zeros((len(groups), n_sphere), dtype=bool)
    for group_id, members in enumerate(groups):
        bool_result[group_id, members] = True
    return bool_result, groups

# 测试示例
if __name__ == "__main__":
    radii = np.array([2, 1, 1, 2, 2, 0.5])
    poss = np.array([[7, 7, 7], [7.5, 8.5, 6], [0, 0, 0], [-1, -2, -1], [1, 1, 1], [2, 1, 3]])
    bool_arr, index_groups = sphere_connectivity_group(poss, radii)
    print("布尔数组输出：\n", bool_arr)
    print("索引分组输出：\n", index_groups)

性能说明 #

• 对于10万个均匀分布的球体，该方案全流程耗时通常在数秒级别，内存开销可控。
• 如果球体分布极不均匀（存在大量紧密堆积的区域），可以将KDTree替换为分层空间哈希结构，进一步提升粗筛效率。
• 不推荐在该场景下使用networkx：其纯Python实现的图结构内存开销大，连通分量遍历效率比原生并查集低1~2个数量级，不适合十万级以上的数据规模。

内容的提问来源于stack exchange，提问作者Ali_Sh