AtCoder Beginner Contest 350 D

https://atcoder.jp/contests/abc350/tasks/abc350_d

連結成分に分解して、個々のグラフが完全グラフにするのに何本エッジが要るかを調べます。

// New Friends
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Graph ////////////////////

use std::collections::HashMap;

type Node = usize;
type Edge = (Node, Node);

fn read_edge() -> Edge {
    let v = read_vec::<Node>();
    (v[0]-1, v[1]-1)
}

struct Graph {
    g: Vec<Vec<Node>>
}

impl Graph {
    fn len(&self) -> usize {
        self.g.len()
    }
    
    fn num_edges(&self) -> usize {
        self.g.iter().map(|vs| vs.len()).sum::<usize>() / 2
    }
    
    fn divide_into_connected(&self) -> Vec<Graph> {
        let N = self.g.len();
        let mut u = UnionFind::new(N);
        for (n, v) in self.g.iter().enumerate() {
            for &m in v.iter() {
                u.join(n, m)
            }
        }
        
        let mut m: HashMap<Node, Vec<Node>> = HashMap::new();
        for v in 0..N {
            let r = u.root(v);
            let e = m.entry(r).or_insert(vec![]);
            (*e).push(v)
        }
        
        m.into_iter().
            map(|(_, vs)| Graph { g: vs.into_iter().map(|v| self.g[v].to_vec()).
                                                collect::<Vec<Vec<Node>>>() }).
            collect::<Vec<Graph>>()
    }
    
    fn create(N: usize, edges: Vec<Edge>) -> Graph {
        let mut g: Vec<Vec<Node>> = (0..N).map(|_| vec![]).collect();
        for (u, v) in edges.into_iter() {
            g[u].push(v);
            g[v].push(u)
        }
        Graph { g }
    }
}


//////////////////// UnionFind ////////////////////

use std::cmp::max;

struct UnionFind {
    parents: Vec<Node>,
    heights: Vec<i32>,
}

impl UnionFind {
    fn new(N: usize) -> UnionFind {
        let parents: Vec<Node> = (0..N).collect();
        let heights: Vec<i32> = vec![1; N];
        UnionFind { parents, heights }
    }
    
    fn is_root(&self, v: Node) -> bool {
        self.parents[v] == v
    }
    
    fn join(&mut self, u: Node, v: Node) {
        let r1 = self.root(u);
        let r2 = self.root(v);
        if r1 == r2 {
            return  // ここには来ないはず
        }
        
        let h1 = self.heights[r1];
        let h2 = self.heights[r2];
        if h1 <= h2 {   // r2にr1がぶら下がる
            self.parents[r1] = r2;
            self.heights[r2] = max(self.heights[r2], self.heights[r1]+1);
        }
        else {
            self.parents[r2] = r1;
            self.heights[r1] = max(self.heights[r1], self.heights[r2]+1);
        }
    }
    
    fn root(&self, v0: Node) -> Node {
        if self.is_root(v0) {
            v0
        }
        else {
            let p = self.parents[v0];
            self.root(p)
        }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (usize, Vec<Edge>) {
    let v = read_vec();
    let (N, M) = (v[0], v[1]);
    let edges: Vec<Edge> = (0..M).map(|_| read_edge()).collect();
    (N, edges)
}

fn F(N: usize, edges: Vec<Edge>) -> usize {
    let graph = Graph::create(N, edges);
    let subgraphs = graph.divide_into_connected();
    subgraphs.iter().map(|g| (g.len() * (g.len()-1) / 2 - g.num_edges())).
                                                            sum::<usize>()
}

fn main() {
    let (N, edges) = read_input();
    println!("{}", F(N, edges))
}