import Foundation

/// 経路圧縮 + ランク併合の Union-Find。類似ペアから連結成分（グループ）を作るのに使う。
struct UnionFind {
    private var parent: [Int]
    private var rank: [Int]

    init(count: Int) {
        parent = Array(0..<count)
        rank = Array(repeating: 0, count: count)
    }

    mutating func find(_ x: Int) -> Int {
        var root = x
        while parent[root] != root { root = parent[root] }
        // 経路圧縮
        var node = x
        while parent[node] != root {
            let next = parent[node]
            parent[node] = root
            node = next
        }
        return root
    }

    @discardableResult
    mutating func union(_ a: Int, _ b: Int) -> Bool {
        let ra = find(a)
        let rb = find(b)
        guard ra != rb else { return false }
        if rank[ra] < rank[rb] {
            parent[ra] = rb
        } else if rank[ra] > rank[rb] {
            parent[rb] = ra
        } else {
            parent[rb] = ra
            rank[ra] += 1
        }
        return true
    }

    /// 連結成分を、根のインデックスをキーにしてまとめて返す。
    mutating func components() -> [[Int]] {
        var buckets: [Int: [Int]] = [:]
        for i in parent.indices {
            buckets[find(i), default: []].append(i)
        }
        return Array(buckets.values)
    }
}