基于Union-Find的有向图树根查找算法错误排查

我想实现一个基于SwiftGraph的算法，从任意图（不一定是树）中查找可行的树根，找到则返回其在顶点列表中的索引，否则返回nil。我实现了带路径压缩的按秩合并Union-Find，并为Graph扩展了findTreeRoot方法，但该方法存在返回错误顶点索引的问题。以下是我的实现代码及失败测试用例，请求排查问题所在：

带路径压缩的按秩合并Union-Find： #

fileprivate class UnionFind {
    private var parent: [Int]
    private var rank: [Int]
    
    init(elements: [Int]) {
        self.parent = [Int].init(repeating: 0, count: elements.count)
        self.rank = [Int].init(repeating: 0, count: elements.count)
        
        self.makeSet(elements)
    }

    private final func makeSet(_ vertices: [Int]) {
        for i in 0..<vertices.count {
            parent[i] = i
            rank[i] = 0
        }
    }

    func find(_ x: Int) -> Int {
        if parent[x] != x {
            parent[x] = find(parent[x]) // Path compression
        }
        return parent[x]
    }

    func union(_ x: Int, _ y: Int) {
        let xRoot = find(x)
        let yRoot = find(y)
        
        if xRoot == yRoot {
            return
        }
        
        if rank[xRoot] < rank[yRoot] {
            parent[xRoot] = yRoot
        } else if rank[xRoot] > rank[yRoot] {
            parent[yRoot] = xRoot
        } else {
            parent[yRoot] = xRoot
            rank[xRoot] += 1
        }
    }
}

SwiftGraph扩展： #

extension Graph {
    /// If this graph is a valid tree, this algorihtm finds and returns the index of the root vertex in `self.vertices`.
    ///
    /// May require further testing.
    ///
    /// - Returns: The index of the root of the tree, if this `self` is a tree, `nil` otherwise. If the graph is a degenerate tree,
    /// i.e. an empty tree, this function returns nil.
    ///
    /// - Complexity: **Time:** O(V + E⋅α(V)) where α(V) is the inverse Ackermann function, that grows extremely slowly (it is considered
    /// constant for any common practical application). **Memory:** O(V) to create the sets for Union-Find by rank with path compression.
    public func findTreeRoot() -> Int? {
        let unionFind = UnionFind(elements: [Int](self.vertices.indices))
        
        for vertex in 0..<self.vertices.count {
            for neighbor in self.edges[vertex] {
                let rootVertex = unionFind.find(vertex)
                let rootNeighbor = unionFind.find(neighbor.v)
                
                if rootVertex == rootNeighbor {
                    return nil
                }
                
                unionFind.union(rootVertex, rootNeighbor)
            }
        }

        var roots = Set<Int>()
        for vertex in 0..<self.vertices.count {
            roots.insert(unionFind.find(vertex))
        }
    
        if roots.count == 1 {
            return roots.first
        } else {
            return nil
        }
    }
}

失败测试用例： #

let originalGraph = WeightedGraph<String, Float>()
        
for component in ["topbar", "secondary topbar", "gallery", "toolbar", "captions"] {
    let _ = originalGraph.addVertex(component)
}
                
originalGraph.addEdge(fromIndex: 0, toIndex: 1, weight: 1.2, directed: true)
originalGraph.addEdge(fromIndex: 0, toIndex: 2, weight: 0.5, directed: true)
originalGraph.addEdge(fromIndex: 0, toIndex: 3, weight: 1.1, directed: true)
originalGraph.addEdge(fromIndex: 0, toIndex: 4, weight: 1.3, directed: true)
        
originalGraph.addEdge(fromIndex: 1, toIndex: 0, weight: 1.1, directed: true)
originalGraph.addEdge(fromIndex: 1, toIndex: 2, weight: 0.5, directed: true)
originalGraph.addEdge(fromIndex: 1, toIndex: 3, weight: 1.5, directed: true)
originalGraph.addEdge(fromIndex: 1, toIndex: 4, weight: 2.2, directed: true)

originalGraph.addEdge(fromIndex: 3, toIndex: 2, weight: 1.0, directed: true)
originalGraph.addEdge(fromIndex: 3, toIndex: 4, weight: 1.0, directed: true)
        
assert(originalGraph.findTreeRoot() == nil)

let msaOfTopbar = try originalGraph.msa(root: 1)
        
let msaAsGraph = WeightedGraph<String, Float>(vertices: originalGraph.vertices)
for edge in msaOfTopbar {
    msaAsGraph.addEdge(edge, directed: true)
}
        
assert(msaAsGraph.findTreeRoot() == 1) // assertion failed: 0 != 1

核心问题分析 #

1. 有向图被当作无向图处理：你的findTreeRoot方法使用Union-Find合并边的两个顶点，但完全忽略了边的方向。有向树的根是入度为0的节点，且所有其他节点都有且仅有一条入边，而Union-Find的合并逻辑无法区分这种有向结构，导致最终返回的Union-Find根节点和实际的有向树根不匹配。

2. 错误的树判定逻辑：当前代码只能判断无向图是否是一棵无环连通树，无法处理有向树的场景。在测试用例中，msaAsGraph是一棵以1为根的有向树，但Union-Find合并0和1时，因初始秩相同，会将1的父节点设为0，后续所有节点最终都会合并到0的集合，所以返回0而非正确的根1。

修正后的实现 #

针对有向树的判定，应该通过入度统计+遍历验证的方式实现，这样才能准确找到有向树的根：

extension Graph {
    public func findTreeRoot() -> Int? {
        guard !vertices.isEmpty else { return nil }
        
        // 1. 计算每个节点的入度
        var inDegree = [Int](repeating: 0, count: vertices.count)
        for vertex in 0..<vertices.count {
            for neighbor in edges[vertex] {
                inDegree[neighbor.v] += 1
            }
        }
        
        // 2. 查找入度为0的节点，有向树必须恰好存在一个这样的节点
        var rootCandidates = [Int]()
        for (index, degree) in inDegree.enumerated() {
            if degree == 0 {
                rootCandidates.append(index)
            }
        }
        guard rootCandidates.count == 1 else { return nil }
        let candidateRoot = rootCandidates[0]
        
        // 3. DFS验证：从候选根出发能遍历所有节点，且无环
        var visited = [Bool](repeating: false, count: vertices.count)
        var visitedCount = 0
        var hasCycle = false
        
        func dfs(_ node: Int, _ parent: Int) {
            if visited[node] {
                hasCycle = true
                return
            }
            visited[node] = true
            visitedCount += 1
            
            for neighbor in edges[node] {
                let nextNode = neighbor.v
                dfs(nextNode, node)
                if hasCycle { return }
            }
        }
        
        dfs(candidateRoot, -1)
        
        // 4. 确认无环且所有节点都被访问到
        return !hasCycle && visitedCount == vertices.count ? candidateRoot : nil
    }
}

修正逻辑说明 #

1. 入度统计：有向树的根唯一特征是入度为0，其他节点入度均为1，这一步先筛选出唯一的候选根。
2. DFS验证：从候选根出发遍历整个图，确保所有节点都能被访问到（连通性），同时检测是否存在环（有环则不是树）。
3. 结果判定：只有满足无环、连通、且唯一入度为0的节点存在时，才返回该节点作为根，否则返回nil。

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