Union-Find求解最大组件尺寸时每次运行结果不一致问题排查

问题描述 #

使用并查集求解「按公共因子计算最大组件尺寸」问题时，同一输入在Xcode Playground中每次运行结果不同。例如输入数组[4, 6, 15, 35]，调用largestComponentSize函数后，结果可能为2、3或4。

原代码 #

class Solution {
    var uf = UnionFind()
    
    func largestComponentSize(_ nums: [Int]) -> Int {
        var maxNum:Int = 0
        var numFactorMap = [Int:Int]()
        var factorAdded = Set<Int>()
        for num in nums {
            var pFactors = getPrimeFactors(num)
            numFactorMap[num] = pFactors[0]
            for (i, val) in pFactors.enumerated() {
                if !factorAdded.contains(val) {
                    uf.addSet(val)
                    factorAdded.insert(val)
                }
                if i > 0 {
                    uf.union(pFactors[i-1], val)
                }
            }
        }
        
        var groupCountMap = [Int:Int]()
        
        for num in nums {
            var groupId = uf.find(numFactorMap[num]!)!
            if groupCountMap.keys.contains(groupId) {
                groupCountMap[groupId]! += 1
            } else {
                groupCountMap[groupId] = 1
            }
            maxNum = max(maxNum, groupCountMap[groupId]!)
        }
        
        return maxNum
    }

    func getPrimeFactors(_ num: Int) -> [Int] {
        var ans:Set<Int> = []

        if num == 1 {
            return []
        }
        
        var crrNum = num
        var deno = 2
        while crrNum >= deno {
            if crrNum % deno == 0 {
                ans.insert(deno)
                crrNum = crrNum / deno
            } else {
                deno = deno + 1
            }
        }
        
        return Array(ans)
    }
    
    class UnionFind {
        var index = [Int: Int]()
        var parent: [Int]
        var size: [Int]
        
        init() {
            parent = []
            size = []
        }
        
        func addSet(_ ele: Int) {
            index[ele] = parent.count
            parent.append(parent.count)
            size.append(1)
        }
        
        func getSetSize(_ ele: Int) -> Int {
            if let found = find(ele) {
                return size[found]
            }
            return 0
        }
        
        func find(_ ele: Int) -> Int? {
            if let indexOfEle = index[ele] {
                if parent[indexOfEle] == indexOfEle {
                    return indexOfEle
                } else {
                    if let found = find(parent[indexOfEle]) {
                        parent[indexOfEle] = found
                    }
                    return parent[indexOfEle]
                }
            } else {
                return nil //never come here
            }
        }
        
        func union(_ first: Int, _ second: Int) {
            guard let indexOfFirst = index[first], let indexOfSecond = index[second] else {
                return
            }
            
            if parent[indexOfFirst] == parent[indexOfSecond] {
                return
            }
            
            var indexOfLarger = indexOfFirst
            var indexOfSmaller = indexOfSecond
            if size[indexOfFirst] < size[indexOfSecond] {
                indexOfLarger = indexOfSecond
                indexOfSmaller = indexOfFirst
            }
            
            parent[indexOfSmaller] = indexOfLarger
            size[indexOfLarger] += size[indexOfSmaller]
            
            return
        }
    }
}

var sol = Solution()
var nums = [4, 6, 15, 35]
var ans = sol.largestComponentSize(nums)

问题原因分析 #

1. 质因子集合无序导致映射不稳定
   getPrimeFactors方法使用Set存储质因子，而Set是无序集合，转换为数组后pFactors[0]的取值每次可能不同。例如6的质因子是2和3，转换为数组时可能是[2,3]或[3,2]，导致numFactorMap[6]的取值不确定，后续分组时会错误关联到不同的质因子组。

2. 并查集Union方法判断逻辑错误
   UnionFind的union方法中，直接比较parent[indexOfFirst]和parent[indexOfSecond]来判断是否属于同一集合，这是错误的。因为parent[indexOfFirst]可能不是根节点，只是路径中的中间节点，会导致本该合并的集合未合并，或错误判断集合归属，最终导致组件尺寸计算错误。

修复方案 #

1. 稳定质因子数组的顺序 #

将getPrimeFactors中的Set转换为有序数组，确保每次返回的质因子顺序一致：

func getPrimeFactors(_ num: Int) -> [Int] {
    var ans = Set<Int>()

    if num == 1 {
        return []
    }
    
    var crrNum = num
    var deno = 2
    while crrNum >= deno {
        if crrNum % deno == 0 {
            ans.insert(deno)
            crrNum = crrNum / deno
        } else {
            deno += 1
        }
    }
    
    // 将集合转换为有序数组，确保每次返回顺序一致
    return ans.sorted()
}

2. 修正并查集Union方法的判断逻辑 #

在union方法中，先通过find获取两个元素的根节点，再判断根节点是否相同，确保合并逻辑正确：

func union(_ first: Int, _ second: Int) {
    guard let rootFirst = find(first), let rootSecond = find(second) else {
        return
    }
    
    // 根节点相同则无需合并
    if rootFirst == rootSecond {
        return
    }
    
    // 按大小合并，将小集合合并到大集合中
    var indexOfLarger = rootFirst
    var indexOfSmaller = rootSecond
    if size[rootFirst] < size[rootSecond] {
        indexOfLarger = rootSecond
        indexOfSmaller = rootFirst
    }
    
    parent[indexOfSmaller] = indexOfLarger
    size[indexOfLarger] += size[indexOfSmaller]
}

修复后效果 #

修改后，输入[4,6,15,35]的稳定结果应为4，因为所有数字通过公共因子连通：4(2)→6(2,3)→15(3,5)→35(5,7)，所有4个数字属于同一个组件。

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