Created
March 16, 2018 06:46
-
-
Save pcoutin/ec41be47b9500475160fd47e21c478e5 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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)); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment