https://atcoder.jp/contests/abc328/tasks/abc328_f
これもUnion-Findですね。親との差をメモしておくと、元々繋がっている時に矛盾するかチェックできます。
// Good Set Query #![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() } //////////////////// UnionFind //////////////////// type Node = usize; type Dist = i64; type Edge = (Node, Node, Dist); struct UnionFind { parents: Vec<Node>, heights: Vec<i32>, dists: Vec<i64> // from parent } impl UnionFind { fn new(N: usize) -> UnionFind { let parents: Vec<Node> = (0..N).collect(); let heights: Vec<i32> = vec![1; N]; let dists: Vec<Dist> = vec![0; N]; UnionFind { parents, heights, dists } } fn is_root(&self, v: Node) -> bool { self.parents[v] == v } fn join(&mut self, edge: &Edge) { let (r1, d1) = self.root(edge.0); let (r2, d2) = self.root(edge.1); 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); self.dists[r1] = d2 - d1 - edge.2 } else { self.parents[r2] = r1; self.heights[r1] = max(self.heights[r1], self.heights[r2]+1); self.dists[r2] = d1 - d2 + edge.2 } } fn root(&self, v0: Node) -> (Node, Dist) { if self.is_root(v0) { (v0, 0) } else { let p = self.parents[v0]; let (r, d) = self.root(p); (r, self.dists[v0] + d) } } fn is_connected(&self, v1: Node, v2: Node) -> bool { let (r1, _) = self.root(v1); let (r2, _) = self.root(v2); r1 == r2 } fn is_valid(&self, edge: Edge) -> bool { let (_r1, d1) = self.root(edge.0); let (_r2, d2) = self.root(edge.1); d2 - d1 == edge.2 } } //////////////////// process //////////////////// fn read_input() -> (usize, usize) { let v = read_vec(); let (N, Q) = (v[0], v[1]); (N, Q) } fn read_edge() -> Edge { let v: Vec<Dist> = read_vec(); let a: Node = (v[0] - 1) as Node; let b: Node = (v[1] - 1) as Node; let d = v[2]; (a, b, d) } fn F(N: usize, Q: usize) { let mut nodes: Vec<usize> = vec![]; let mut uf = UnionFind::new(N); for i in 0..Q { let edge = read_edge(); if uf.is_connected(edge.0, edge.1) { if uf.is_valid(edge) { nodes.push(i + 1) } } else { uf.join(&edge); nodes.push(i + 1) } } let ss: Vec<String> = nodes.iter().map(|i| i.to_string()).collect(); println!("{}", ss.join(" ")) } fn main() { let (N, Q) = read_input(); F(N, Q) }