51 lines
1.4 KiB
Swift
51 lines
1.4 KiB
Swift
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)
|
|
}
|
|
}
|