基于阈值对节点关系对称矩阵中的节点进行分组的函数开发需求

问题描述 #

给定N个索引为0到N-1的节点，它们的连接权重由一个N×N的对称矩阵表示（权重范围0~1，对角线元素恒为1）。需要实现一个函数，输入阈值和该矩阵，将节点分组：若两个节点间的边权重超过阈值，则这两个节点归为同一组（同一连通分量）。

示例权重矩阵：

import numpy as np

weight_matrix = np.array([[1.0, 0.2, 0.3, 0.4],
                          [0.2, 1.0, 0.7, 0.5],
                          [0.3, 0.7, 1.0, 0.2],
                          [0.4, 0.5, 0.2, 1.0]])

不同阈值的预期分组结果：

• 阈值=0.6：[[0], [1,2], [3]]（节点1和2的边权重0.7>0.6，归为一组；其余节点无符合条件的连接）
• 阈值=0.8：[[0],[1],[2],[3]]（所有非对角线边权重都≤0.8，每个节点单独一组）
• 阈值=0.1：[[0,1,2,3]]（所有边权重都>0.1，所有节点归为一组）

解决方案：并查集（Union-Find）算法 #

并查集是处理连通分量问题的高效方案，适合这类节点分组场景。我们遍历矩阵中所有i<j的节点对（利用矩阵对称性避免重复处理），当权重超过阈值时合并两个节点的集合，最后收集所有连通分量即可。

代码实现 #

import numpy as np

def group_nodes(threshold: float, weight_matrix: np.ndarray) -> list[list[int]]:
    n = weight_matrix.shape[0]
    # 初始化并查集：每个节点的父节点指向自己
    parent = list(range(n))
    
    def find(u: int) -> int:
        # 路径压缩优化，快速定位根节点
        if parent[u] != u:
            parent[u] = find(parent[u])
        return parent[u]
    
    def union(u: int, v: int):
        # 合并两个节点所在的集合
        root_u = find(u)
        root_v = find(v)
        if root_u != root_v:
            parent[root_v] = root_u
    
    # 遍历所有非重复节点对
    for i in range(n):
        for j in range(i+1, n):
            if weight_matrix[i][j] > threshold:
                union(i, j)
    
    # 按根节点收集连通分量
    groups = {}
    for node in range(n):
        root = find(node)
        groups.setdefault(root, []).append(node)
    
    # 返回排序后的分组（保证结果整洁一致）
    return [sorted(group) for group in groups.values()]

测试示例 #

# 测试示例矩阵
weight_matrix = np.array([[1.0, 0.2, 0.3, 0.4],
                          [0.2, 1.0, 0.7, 0.5],
                          [0.3, 0.7, 1.0, 0.2],
                          [0.4, 0.5, 0.2, 1.0]])

print(group_nodes(0.6, weight_matrix))  # 输出: [[0], [1, 2], [3]]
print(group_nodes(0.8, weight_matrix))  # 输出: [[0], [1], [2], [3]]
print(group_nodes(0.1, weight_matrix))  # 输出: [[0, 1, 2, 3]]

代码说明 #

1. 并查集初始化：每个节点初始为独立集合，父节点指向自身。
2. find函数：通过路径压缩优化，将节点直接指向根节点，降低后续查找的时间成本。
3. union函数：将两个节点的集合合并，通过修改根节点的父节点实现。
4. 遍历节点对：仅处理i<j的节点对，利用矩阵对称性避免重复操作，权重超过阈值时执行合并。
5. 收集分组：按根节点归类所有节点，返回排序后的分组列表（排序为可选操作，用于对齐示例结果）。

注意事项 #

• 输入矩阵需为对称矩阵，函数无需依赖对角线元素（遍历逻辑不会处理对角线）。
• 权重范围0~1，函数对超出范围的权重也能处理，但需符合业务逻辑。
• 时间复杂度接近O(N²)（并查集操作近乎O(1)），适合中等规模的节点数量；若节点量极大，可考虑仅存储符合条件的边来优化遍历。

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