AtCoder Beginner Contest 328 F

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