haar_lib/ds/
potential_unionfind.rs

1//! ポテンシャル付きUnionfind
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/unionfind_with_potential>
5//! - <https://judge.yosupo.jp/problem/unionfind_with_potential_non_commutative_group>
6
7use crate::algebra::traits::*;
8use std::cell::{Cell, RefCell};
9
10/// ポテンシャル付きUnionfind
11pub struct PotentialUnionFind<T> {
12    n: usize,
13    count: usize,
14    parent: Vec<Cell<usize>>,
15    depth: Vec<usize>,
16    size: Vec<usize>,
17    potential: RefCell<Vec<T>>,
18    inverse: Option<RefCell<Vec<T>>>,
19}
20
21impl<T> PotentialUnionFind<T>
22where
23    T: AbelianGroup + Clone,
24{
25    /// 大きさ`n`の[`PotentialUnionFind`]を生成する。(ポテンシャルが可換群のとき)
26    pub fn new_commutative(n: usize) -> Self {
27        Self {
28            n,
29            count: n,
30            parent: (0..n).map(Cell::new).collect(),
31            depth: vec![1; n],
32            size: vec![1; n],
33            potential: RefCell::new(vec![T::id(); n]),
34            inverse: None,
35        }
36    }
37}
38
39impl<T> PotentialUnionFind<T>
40where
41    T: Group + Clone,
42{
43    /// 大きさ`n`の[`PotentialUnionFind`]を生成する。(ポテンシャルが可換とは限らない群のとき)
44    pub fn new_non_commutative(n: usize) -> Self {
45        Self {
46            n,
47            count: n,
48            parent: (0..n).map(Cell::new).collect(),
49            depth: vec![1; n],
50            size: vec![1; n],
51            potential: RefCell::new(vec![T::id(); n]),
52            inverse: Some(RefCell::new(vec![T::id(); n])),
53        }
54    }
55
56    /// `i`の属する素集合の根を返す。
57    pub fn root_of(&self, i: usize) -> usize {
58        if self.parent[i].get() == i {
59            return i;
60        }
61        let p = self.parent[i].get();
62        let p = self.root_of(p);
63
64        let mut potential = self.potential.borrow_mut();
65        let t = potential[self.parent[i].get()].clone();
66        potential[i].op_assign_l(t);
67
68        if let Some(inv) = self.inverse.as_ref() {
69            let mut inv = inv.borrow_mut();
70            let t = inv[self.parent[i].get()].clone();
71            inv[i].op_assign_r(t);
72        }
73
74        self.parent[i].set(p);
75        self.parent[i].get()
76    }
77
78    /// `i`のポテンシャル($P(i)$)を返す。
79    pub fn potential_of(&self, i: usize) -> T {
80        self.potential.borrow()[i].clone()
81    }
82
83    /// `i`と`j`が同じ素集合に属するならば`true`を返す。
84    pub fn is_same(&self, i: usize, j: usize) -> bool {
85        self.root_of(i) == self.root_of(j)
86    }
87
88    /// `i`と`j`が同一の素集合に属するとき、ポテンシャルの差($P(i) - P(j)$)を返す。
89    pub fn diff(&self, i: usize, j: usize) -> Option<T> {
90        self.is_same(i, j).then(|| {
91            let pi = self.potential_of(i);
92            if let Some(inv) = self.inverse.as_ref() {
93                inv.borrow()[j].clone().op(pi)
94            } else {
95                self.potential_of(j).inv().op(pi)
96            }
97        })
98    }
99
100    /// `i`の属する素集合と`j`の属する素集合を統合する。
101    /// 統合後は、$P(i) = P(j) + p$を満たす。
102    pub fn merge(&mut self, i: usize, j: usize, p: T) -> usize {
103        let ri = self.root_of(i);
104        let rj = self.root_of(j);
105
106        if ri == rj {
107            return ri;
108        }
109
110        self.count -= 1;
111
112        let mut potential = self.potential.borrow_mut();
113        let (pi, pj) = (potential[i].clone(), potential[j].clone());
114
115        if self.depth[ri] < self.depth[rj] {
116            self.parent[ri].set(rj);
117            self.size[rj] += self.size[ri];
118
119            if let Some(inv) = self.inverse.as_ref() {
120                let mut inv = inv.borrow_mut();
121                potential[ri] = pj.op(p.clone()).op(inv[i].clone());
122                inv[ri] = pi.op(p.inv()).op(inv[j].clone());
123            } else {
124                potential[ri] = pj.op(p).op(pi.inv());
125            }
126
127            j
128        } else {
129            self.parent[rj].set(ri);
130            self.size[ri] += self.size[rj];
131
132            if let Some(inv) = self.inverse.as_ref() {
133                let mut inv = inv.borrow_mut();
134                potential[rj] = pi.op(p.clone().inv()).op(inv[j].clone());
135                inv[rj] = pj.op(p).op(inv[i].clone())
136            } else {
137                potential[rj] = pi.op(p.inv()).op(pj.inv());
138            }
139
140            if self.depth[ri] == self.depth[rj] {
141                self.depth[i] += 1;
142            }
143
144            i
145        }
146    }
147
148    /// `i`の属する素集合の大きさを返す。
149    pub fn size_of(&self, i: usize) -> usize {
150        self.size[self.root_of(i)]
151    }
152
153    /// 素集合の個数を返す。
154    pub fn count_groups(&self) -> usize {
155        self.count
156    }
157
158    /// 素集合をすべて列挙する。
159    pub fn get_groups(&self) -> Vec<Vec<usize>> {
160        let mut ret = vec![vec![]; self.n];
161
162        for i in 0..self.n {
163            ret[self.root_of(i)].push(i);
164        }
165
166        ret.into_iter().filter(|x| !x.is_empty()).collect()
167    }
168}
haar_lib/ds/potential_unionfind.rs

haar_lib/ds/
potential_unionfind.rs
