haar_lib/ds/
rollbackable_unionfind.rs

1//! ロールバック可能Unionfind
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/persistent_unionfind>
5
6enum History {
7    Nop,
8    Merge(usize, usize, usize),
9}
10
11/// ロールバック可能Unionfind
12pub struct RollbackableUnionFind {
13    parent: Vec<Option<usize>>,
14    size: Vec<usize>,
15    history: Vec<History>,
16}
17
18impl RollbackableUnionFind {
19    /// `RollbackableUnionFind`を生成する
20    pub fn new(n: usize) -> Self {
21        Self {
22            parent: vec![None; n],
23            size: vec![1; n],
24            history: vec![],
25        }
26    }
27
28    /// `i`の属する素集合の根を返す
29    pub fn root_of(&self, i: usize) -> usize {
30        if let Some(p) = self.parent[i] {
31            self.root_of(p)
32        } else {
33            i
34        }
35    }
36
37    /// `i`と`j`が同じ素集合に属するかを判定する
38    pub fn is_same(&self, i: usize, j: usize) -> bool {
39        self.root_of(i) == self.root_of(j)
40    }
41
42    /// `i`の属する素集合と`j`の属する素集合を統合する
43    pub fn merge(&mut self, i: usize, j: usize) -> usize {
44        let ri = self.root_of(i);
45        let rj = self.root_of(j);
46
47        if ri == rj {
48            self.history.push(History::Nop);
49            return ri;
50        }
51
52        if self.size[ri] < self.size[rj] {
53            self.history.push(History::Merge(ri, rj, self.size[rj]));
54            self.parent[ri] = Some(rj);
55            self.size[rj] += self.size[ri];
56            rj
57        } else {
58            self.history.push(History::Merge(rj, ri, self.size[ri]));
59            self.parent[rj] = Some(ri);
60            self.size[ri] += self.size[rj];
61            ri
62        }
63    }
64
65    /// 直前の`merge`操作を巻き戻す
66    pub fn rollback(&mut self) -> bool {
67        match self.history.pop() {
68            Some(History::Nop) => true,
69            Some(History::Merge(c, p, s)) => {
70                self.parent[c] = None;
71                self.size[p] = s;
72                true
73            }
74            None => false,
75        }
76    }
77
78    /// `i`の属する素集合の大きさを返す
79    pub fn size_of(&self, i: usize) -> usize {
80        self.size[self.root_of(i)]
81    }
82}
haar_lib/ds/rollbackable_unionfind.rs

haar_lib/ds/
rollbackable_unionfind.rs
