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)) }