Submission #18834680

---

Source Code Expand

use std::collections::BTreeMap;

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());
    let n: usize = sc.read();
    let q: usize = sc.read();
    let c: Vec<usize> = sc.vec(n);
    let mut uf = UnionFind::new(&c);
    for _ in 0..q {
        let t: usize = sc.read();
        if t == 1 {
            let a = sc.read::<usize>() - 1;
            let b = sc.read::<usize>() - 1;
            uf.unite(a, b);
        } else {
            let x = sc.read::<usize>() - 1;
            let y = sc.read::<usize>();
            let root = uf.find(x);
            let ans = uf.sizes[root]
                .as_ref()
                .unwrap()
                .get(&y)
                .cloned()
                .unwrap_or(0);
            sc.write(ans);
            sc.write('\n');
        }
    }
}

pub struct UnionFind {
    parent: Vec<usize>,
    sizes: Vec<Option<BTreeMap<usize, usize>>>,
    size: usize,
}

impl UnionFind {
    pub fn new(c: &[usize]) -> UnionFind {
        let n = c.len();
        let mut sizes = vec![Some(BTreeMap::new()); n];
        for i in 0..n {
            sizes[i].as_mut().unwrap().insert(c[i], 1);
        }
        UnionFind {
            parent: (0..n).map(|i| i).collect::<Vec<usize>>(),
            sizes,
            size: n,
        }
    }

    pub fn find(&mut self, x: usize) -> usize {
        if x == self.parent[x] {
            x
        } else {
            let px = self.parent[x];
            self.parent[x] = self.find(px);
            self.parent[x]
        }
    }

    pub fn unite(&mut self, x: usize, y: usize) -> bool {
        let parent_x = self.find(x);
        let parent_y = self.find(y);
        if parent_x == parent_y {
            return false;
        }

        let (large, small) = if self.sizes[parent_x].as_ref().unwrap().len()
            < self.sizes[parent_y].as_ref().unwrap().len()
        {
            (parent_y, parent_x)
        } else {
            (parent_x, parent_y)
        };

        self.parent[small] = large;
        let small_sizes = self.sizes[small].take().unwrap();
        for (color, count) in small_sizes {
            *self.sizes[large]
                .as_mut()
                .unwrap()
                .entry(color)
                .or_insert(0) += count;
        }

        self.size -= 1;
        true
    }
}

pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> Self {
        Self(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write_all(s.to_string().as_bytes()).unwrap();
    }
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}

Submission Info

Judge Result