Merging Communities

  • + 0 comments

    Solution written in Rust:

    use std::io::{self, BufRead};
    
    fn find(parent: &mut Vec<usize>, i: usize) -> usize {
        if parent[i] != i {
            parent[i] = find(parent, parent[i]);
        }
        parent[i]
    }
    
    fn union(parent: &mut Vec<usize>, size: &mut Vec<usize>, a: usize, b: usize) {
        let root_a = find(parent, a);
        let root_b = find(parent, b);
        if root_a != root_b {
            let (big, small) = if size[root_a] >= size[root_b] {
                (root_a, root_b)
            } else {
                (root_b, root_a)
            };
            parent[big] = small;
            size[small] += size[big];
        }
    }
    
    fn main() {
        let stdin = io::stdin();
        let mut lines = stdin.lock().lines().map(|line| line.unwrap());
    
        if let Some(line) = lines.next() {
            let mut input_iter = line.split_whitespace();
            let n: usize = input_iter.next().unwrap().parse().unwrap();
            let q: usize = input_iter.next().unwrap().parse().unwrap();
    
            let mut parent = (0..=n).collect::<Vec<usize>>();
            let mut size = vec![1; n + 1];
    
            for _ in 0..q {
                if let Some(query) = lines.next() {
                    let mut query_iter = query.split_whitespace();
                    let query_type = query_iter.next().unwrap();
                    let values: Vec<usize> = query_iter.map(|val| val.parse().unwrap()).collect();
    
                    if query_type == "M" {
                        let person1 = values[0];
                        let person2 = values[1];
                        union(&mut parent, &mut size, person1, person2);
                    } else if query_type == "Q" {
                        let person = values[0];
                        let root = find(&mut parent, person);
                        println!("{}", size[root]);
                    }
                }
            }
        }
    }