haar_lib/ds/
partially_persistent_unionfind.rs

1//! 部分永続UnionFind
2//!
3//! # Problems
4//! - [AGC 002 D - Stamp Rally](https://atcoder.jp/contests/agc002/tasks/agc002_d)
5//! - [CODE THANKS FESTIVAL 2017 H - Union Sets](https://atcoder.jp/contests/code-thanks-festival-2017-open/tasks/code_thanks_festival_2017_h)
6
7/// 部分永続UnionFind
8pub struct PartiallyPersistentUnionFind {
9    time: usize,
10    p: Vec<Vec<(usize, usize)>>,
11    par: Vec<usize>,
12    rank: Vec<usize>,
13}
14
15/// ある時間での[`PartiallyPersistentUnionFind`]の状態を参照するための構造体。
16pub struct At<'a> {
17    time: usize,
18    p: &'a Vec<Vec<(usize, usize)>>,
19    par: &'a Vec<usize>,
20}
21
22impl PartiallyPersistentUnionFind {
23    /// 大きさ`size`の[`PartiallyPersistentUnionFind`]を生成する。
24    pub fn new(size: usize) -> Self {
25        Self {
26            time: 0,
27            p: vec![vec![(0, 1)]; size],
28            par: (0..size).collect(),
29            rank: vec![1; size],
30        }
31    }
32
33    /// 時刻tでの状態
34    pub fn at(&self, t: usize) -> At {
35        assert!(t <= self.time);
36        At {
37            time: t,
38            p: &self.p,
39            par: &self.par,
40        }
41    }
42
43    /// 最新時点の状態
44    pub fn latest(&self) -> At {
45        self.at(self.time)
46    }
47
48    /// `u`を含む素集合と`v`を含む素集合を融合する。
49    pub fn merge(&mut self, u: usize, v: usize) {
50        let t = self.time;
51        self.time += 1;
52
53        let u = self.at(t).root_of(u);
54        let v = self.at(t).root_of(v);
55
56        if u == v {
57            return;
58        }
59
60        let s = self.p[u].last().unwrap().1 + self.p[v].last().unwrap().1;
61
62        if self.rank[u] < self.rank[v] {
63            self.par[u] = v;
64            self.par[v] = v;
65            self.p[u].push((self.time, v));
66            self.p[v].push((self.time, s));
67        } else {
68            self.par[u] = u;
69            self.par[v] = u;
70            self.p[v].push((self.time, u));
71            self.p[u].push((self.time, s));
72            if self.rank[u] == self.rank[v] {
73                self.rank[u] += 1;
74            }
75        }
76    }
77}
78
79impl At<'_> {
80    /// `i`を含む素集合の代表の値を返す。
81    pub fn root_of(&self, i: usize) -> usize {
82        let &(t, r) = self.p[i].last().unwrap();
83
84        if self.par[i] == i || t == 0 || self.time < t {
85            i
86        } else if self.time == t {
87            r
88        } else {
89            self.root_of(self.par[i])
90        }
91    }
92
93    /// `u`と`v`が同じ素集合に含まれていれば`true`を返す。
94    pub fn is_same(&self, u: usize, v: usize) -> bool {
95        self.root_of(u) == self.root_of(v)
96    }
97
98    /// `u`が属する素集合の大きさを返す。
99    ///
100    /// **Time Complexity** $O(\log t)$
101    pub fn size_of(&self, u: usize) -> usize {
102        let u = self.root_of(u);
103
104        match self.p[u].binary_search(&(self.time + 1, 0)) {
105            Ok(i) | Err(i) => self.p[u][i - 1].1,
106        }
107    }
108}
haar_lib/ds/partially_persistent_unionfind.rs

haar_lib/ds/
partially_persistent_unionfind.rs
