competitive_library/structure/
disjoint_set_union_undo.rs

1//! Union find undo
2use std::collections::{HashMap, HashSet, VecDeque};
3#[derive(Debug, Clone)]
4enum Node {
5    Root(usize, usize),
6    Child(usize),
7}
8/// UnionFind
9/// 経路圧縮を行わないことで undo を可能にする
10#[derive(Clone, Debug)]
11pub struct DisjointSetUnionRollback {
12    uf: Vec<Node>,
13    history: VecDeque<(usize, Node)>,
14    restore_point: Option<usize>,
15}
16
17impl DisjointSetUnionRollback {
18    /// 要素数 n の dsu を構築する
19    #[inline]
20    pub fn new(n: usize) -> DisjointSetUnionRollback {
21        DisjointSetUnionRollback {
22            uf: vec![Node::Root(1, 1); n],
23            history: VecDeque::new(),
24            restore_point: None,
25        }
26    }
27
28    /// 根を取得
29    /// 経路圧縮を行わない
30    #[inline]
31    pub fn root(&self, target: usize) -> usize {
32        match self.uf[target] {
33            Node::Root(_, _) => target,
34            Node::Child(par) => self.root(par),
35        }
36    }
37
38    /// 対象の木をマージ
39    /// 経路圧縮を行わないため変更されるノード数は高々2
40    /// 変更箇所をスタックで保存
41    #[inline]
42    pub fn unite(&mut self, x: usize, y: usize) -> bool {
43        let rx = self.root(x);
44        let ry = self.root(y);
45        if rx == ry {
46            return false;
47        }
48        self.history.push_back((rx, self.uf[rx].clone()));
49        self.history.push_back((ry, self.uf[ry].clone()));
50        let size_x = self.size(rx);
51        let size_y = self.size(ry);
52        let rank_x = self.rank(rx);
53        let rank_y = self.rank(ry);
54        let (i, j) = if rank_x > rank_y { (rx, ry) } else { (ry, rx) };
55        self.uf[i] = Node::Root(
56            size_x + size_y,
57            (rank_x.min(rank_y) + 1).max(rank_x.max(rank_y)),
58        );
59        self.uf[j] = Node::Child(i);
60
61        true
62    }
63
64    /// 同じ木に存在するか
65    #[inline]
66    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
67        self.root(x) == self.root(y)
68    }
69
70    /// 所属する木のサイズ
71    pub fn size(&mut self, x: usize) -> usize {
72        let root = self.root(x);
73        match self.uf[root] {
74            Node::Root(size, _) => size,
75            Node::Child(_) => 1,
76        }
77    }
78    /// 所属する木のランク
79    #[inline]
80    pub fn rank(&mut self, x: usize) -> usize {
81        let root = self.root(x);
82        match self.uf[root] {
83            Node::Root(_, rank) => rank,
84            Node::Child(_) => 1,
85        }
86    }
87
88    /// unite 操作の undo
89    #[inline]
90    pub fn undo(&mut self) {
91        for _ in 0..2 {
92            let (index, node) = self.history.pop_back().unwrap();
93            self.uf[index] = node;
94        }
95    }
96
97    /// 現時点の状態を保存
98    /// 復元には rollback_snapshot
99    #[inline]
100    pub fn snapshot(&mut self) {
101        self.restore_point = Some(self.history.len() >> 1);
102    }
103
104    /// 現時点での保存されている操作回数を返す
105    #[inline]
106    pub fn get_history_length(&self) -> usize {
107        self.history.len() >> 1
108    }
109
110    /// rollback_snapshot で保存された状態へ復元
111    #[inline]
112    pub fn rollback_snapshot(&mut self) {
113        self.rollback(self.restore_point.unwrap());
114    }
115
116    /// 復元
117    /// 任意のタイミングで get_history_length を実行し取得した 値を使用する
118    #[inline]
119    pub fn rollback(&mut self, n: usize) {
120        assert!(self.history.len() >= n << 1);
121
122        while self.history.len() > n << 1 {
123            self.undo();
124        }
125    }
126
127    /// 同じ木に含まれるノードを返す
128    #[inline]
129    pub fn get_same_group(&mut self, x: usize) -> HashSet<usize> {
130        let root = self.root(x);
131        let mut g = HashSet::new();
132        for i in 0..self.uf.len() {
133            if root == self.root(i) {
134                g.insert(i);
135            }
136        }
137        g
138    }
139
140    /// 全ノードを返却
141    #[inline]
142    pub fn get_all_groups(&mut self) -> HashMap<usize, HashSet<usize>> {
143        let mut map: HashMap<usize, HashSet<usize>> = HashMap::new();
144        for i in 0..self.uf.len() {
145            let root = self.root(i);
146
147            map.entry(root).or_default().insert(i);
148        }
149        map
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn test_dsu_rollback() {
159        let mut dsu = DisjointSetUnionRollback::new(6);
160
161        dsu.unite(0, 1);
162        assert!(dsu.is_same(0, 1));
163        dsu.unite(1, 2);
164        assert!(dsu.is_same(0, 2));
165        assert_eq!(dsu.size(0), 3);
166        assert!(!dsu.is_same(0, 3));
167        dsu.snapshot();
168        dsu.unite(0, 3);
169        dsu.unite(3, 4);
170        dsu.unite(4, 5);
171        assert_eq!(dsu.size(5), 6);
172        assert!(dsu.is_same(0, 5));
173        dsu.undo();
174        assert!(!dsu.is_same(0, 5));
175        dsu.rollback_snapshot();
176        assert!(dsu.is_same(0, 2));
177        assert_eq!(dsu.size(0), 3);
178        assert!(!dsu.is_same(0, 3));
179        dsu.rollback(0);
180        assert!(!dsu.is_same(0, 1));
181        assert_eq!(dsu.get_history_length(), 0);
182    }
183}
competitive_library/structure/disjoint_set_union_undo.rs

competitive_library/structure/
disjoint_set_union_undo.rs
