journey_to_the_moon.rs

use std::io::{self, Read};
use std::collections::HashMap;

fn find_set(x: usize, parents: &mut Vec<usize>) -> usize {
    if x != parents[x] {
        parents[x] = find_set(parents[x], parents);
    }
    parents[x]
}

fn union_find(edgelist: &mut Iterator<Item = (usize, usize)>) -> Vec<usize> {
    let (n, _) = edgelist.next().unwrap();
    let mut parents = (0..n).collect();
    let mut rank = vec![0; n];

    for (x, y) in edgelist {
        if find_set(x, &mut parents) != find_set(y, &mut parents) {
            //merge sets
            let px = find_set(x, &mut parents);
            let py = find_set(y, &mut parents);

            if rank[px] > rank[py] {
                parents[py] = px;
            } else {
                parents[px] = py;
            }

            if rank[px] == rank[py] {
                rank[py] += 1;
            }
        }
    }

    let mut count_in_set = HashMap::new();

    for i in 0..n {
        let s = find_set(i, &mut parents);
        let u = match count_in_set.get(&s) {
                    Some(count) => count + 1,
                    None => 1
                };
        count_in_set.insert(s, u);
    }

    count_in_set.values().map(|x| *x as usize).collect()
}

fn combo_sum(counts: Vec<usize>) -> usize {
    let n = counts.len();
    let mut sum = 0;

    for i in 0..n-1 {
        for j in i+1..n {
            sum += counts[i] * counts[j];
        }
    }
    sum
}

fn edge_of_str(s : &str) -> (usize, usize) {
    let mut edge_it = s
        .split_whitespace()
        .map(|x| x.parse::<usize>().unwrap());

    (edge_it.next().unwrap(), edge_it.next().unwrap())
}

fn main() {
    let mut buffer = String::new();
    let stdin = io::stdin();
    let mut handle = stdin.lock();

    handle.read_to_string(&mut buffer)
        .expect("stdin error");

    let mut input_numbers = buffer.lines().map(edge_of_str);

    let counts = union_find(&mut input_numbers);
    println!("{}", combo_sum(counts));
}