Skip to main content

library_checker/data_structure/
unionfind_with_potential.rs

1use competitive::prelude::*;
2use competitive::{
3    algebra::AdditiveOperation, data_structure::PotentializedUnionFind,
4    num::montgomery::MInt998244353 as M,
5};
6
7competitive::define_enum_scan! {
8    enum Query: u8 {
9        0 => Unite { u: usize, v: usize, x: M }
10        1 => Diff { u: usize, v: usize }
11    }
12}
13
14#[verify::library_checker("unionfind_with_potential")]
15pub fn unionfind_with_potential(reader: impl Read, writer: impl Write) {
16    prepare_io!(reader, writer);
17    sc!(n, q);
18    let mut uf = PotentializedUnionFind::<AdditiveOperation<M>>::new(n);
19    for _ in 0..q {
20        sc!(query: Query);
21        match query {
22            Query::Unite { u, v, x } => {
23                if let Some(diff) = uf.difference(u, v) {
24                    pp!((diff == x) as u8);
25                } else {
26                    uf.unite_with(u, v, x);
27                    pp!("1");
28                }
29            }
30            Query::Diff { u, v } => {
31                if let Some(diff) = uf.difference(u, v) {
32                    pp!(diff);
33                } else {
34                    pp!("-1");
35                }
36            }
37        }
38    }
39}