AtCoder Beginner Contest 352 E

https://atcoder.jp/contests/abc352/tasks/abc352_e

Kruscal法を使うだけです。

// Clique Connect
#![allow(non_snake_case)]

use std::cmp::max;


//////////////////// 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 ////////////////////

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


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

struct UnionFind {
    parents: Vec<Node>,
    heights: Vec<i32>,
    num_edges: usize,
    sum_weights: Weight
}

impl UnionFind {
    fn new(N: usize) -> UnionFind {
        let parents: Vec<Node> = (0..N).collect();
        let heights: Vec<i32> = vec![1; N];
        UnionFind { parents, heights, num_edges: 0, sum_weights: 0 }
    }
    
    fn join(&mut self, edge: &Edge) -> (Node, Node, Node, i32) {
        let r1 = self.root(edge.0);
        let r2 = self.root(edge.1);
        self.num_edges += 1;
        self.sum_weights += edge.2;
        if r1 == r2 {
            return (0, 0, 0, 0) // ここには来ないはず
        }
        
        let h1 = self.heights[r1];
        let h2 = self.heights[r2];
        if h1 <= h2 {   // r2にr1がぶら下がる
            let ret = (r1, self.parents[r1], r2, self.heights[r2]);
            self.parents[r1] = r2;
            self.heights[r2] = max(self.heights[r2], self.heights[r1]+1);
            ret
        }
        else {
            let ret = (r2, self.parents[r2], r1, self.heights[r1]);
            self.parents[r2] = r1;
            self.heights[r1] = max(self.heights[r1], self.heights[r2]+1);
            ret
        }
    }
    
    fn remove(&mut self, r1: Node, p1: Node, r2: Node, h2: i32, w: Weight) {
        self.parents[r1] = p1;
        self.heights[r2] = h2;
        self.num_edges -= 1;
        self.sum_weights -= w
    }
    
    fn root(&self, v0: Node) -> Node {
        let mut v = v0;
        while self.parents[v] != v {
            v = self.parents[v]
        }
        v
    }
    
    fn is_connected(&self, v1: Node, v2: Node) -> bool {
        self.root(v1) == self.root(v2)
    }
}


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

type Group = (Weight, Vec<Node>);

fn read_group() -> Group {
    let v = read_vec();
    let C = v[1];
    let A: Vec<Node> = read_vec::<Node>().into_iter().
                                    map(|a| a - 1).collect();
    (C, A)
}

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

fn F(N: usize, mut groups: Vec<Group>) -> Weight {
    let mut uf = UnionFind::new(N);
    groups.sort_by_key(|&(C, _)| C);
    for (C, A) in groups.into_iter() {
        for i in 0..A.len()-1 {
            if !uf.is_connected(A[i], A[i+1]) {
                uf.join(&(A[i], A[i+1], C));
            }
        }
    }
    
    if uf.num_edges == N - 1 {
        uf.sum_weights
    }
    else {
        -1
    }
}

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