Codeforces/1850/src/bin/h.rs

#![allow(unused_imports)]

// region ‹use›
use self::input::*;
use self::math::*;
use self::output::*;
use std::cmp::{max, min};
use std::collections::HashMap;
// endregion ‹use›

#[derive(Debug, Clone, PartialEq, Eq)]
enum Color {
    White,
    Gray,
    Black,
}

fn dfs(
    g: &[Vec<(usize, i64)>],
    positions: &mut Vec<i64>,
    state: &mut Vec<Color>,
    u: usize,
    parent: usize,
) -> bool {
    state[u] = Color::Gray;

    for (v, d) in &g[u] {
        if *v == parent {
            continue;
        }

        if state[*v] != Color::White {
            if positions[*v] != positions[u] + d {
                // we have already assigned the soldier and it doesn't match
                return false;
            } else {
                continue;
            }
        }

        positions[*v] = positions[u] + d;
        if !dfs(g, positions, state, *v, u) {
            return false;
        }
    }

    state[u] = Color::Black;
    true
}

fn can_arrange(n: usize, conditions: Vec<(usize, usize, i64)>) -> bool {
    let mut graph = vec![vec![]; n];

    // construct graph
    for (u, v, d) in conditions {
        graph[u - 1].push((v - 1, d));
        graph[v - 1].push((u - 1, -d));
    }

    // run DFS
    let mut positions = vec![0; n];
    let mut state: Vec<Color> = vec![Color::White; n];
    for start in 0..n {
        if state[start] != Color::White {
            continue;
        }

        if !dfs(&graph, &mut positions, &mut state, start, n) {
            return false;
        }
    }

    true
}

fn solve(s: &mut Scanner) {
    let soldiers = s.next::<usize>();
    let n = s.next::<usize>();
    let conditions: Vec<(usize, usize, i64)> = (0..n)
        .map(|_| {
            let a = s.next();
            let b = s.next();
            let d = s.next();

            (a, b, d)
        })
        .collect();

    yesno(can_arrange(soldiers, conditions));
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_1() {
        assert!(can_arrange(5, vec![(1, 2, 2), (2, 3, 4), (4, 2, -6)]));
    }
    #[test]
    fn example_2() {
        assert!(!can_arrange(
            6,
            vec![(1, 2, 2), (2, 3, 4), (4, 2, -6), (5, 4, 4), (3, 5, 100)]
        ));
    }
    #[test]
    fn example_3() {
        assert!(!can_arrange(2, vec![(1, 2, 5), (1, 2, 4)]));
    }
    #[test]
    fn example_4() {
        assert!(can_arrange(4, vec![(1, 2, 3)]));
    }
    #[test]
    fn regression_1() {
        assert!(can_arrange(
            4,
            vec![(3, 1, 3), (4, 2, 0), (1, 3, -3), (2, 3, -4)]
        ))
    }
    #[test]
    fn regression_2() {
        assert!(can_arrange(3, vec![(1, 2, 8), (2, 3, -1), (3, 1, -7)]))
    }
    #[test]
    fn regression_3() {
        assert!(!can_arrange(
            9,
            vec![
                (1, 2, 536870912),
                (2, 3, 536870912),
                (3, 4, 536870912),
                (4, 5, 536870912),
                (5, 6, 536870912),
                (6, 7, 536870912),
                (7, 8, 536870912),
                (8, 9, 536870912),
                (1, 9, 0)
            ]
        ));
    }
}

// region runner
const SINGLE_TEST: bool = false;
fn main() {
    let mut s = Scanner::new();

    if SINGLE_TEST {
        solve(&mut s)
    } else {
        let n = s.next::<usize>();
        for _ in 0..n {
            solve(&mut s)
        }
    }
}
// endregion runner

#[allow(dead_code)]
mod math {
    const MOD: i64 = 1_000_000_007;

    pub fn add(a: i64, b: i64) -> i64 {
        (a + b) % MOD
    }

    pub fn sub(a: i64, b: i64) -> i64 {
        ((a - b) % MOD + MOD) % MOD
    }

    pub fn mul(a: i64, b: i64) -> i64 {
        (a * b) % MOD
    }

    pub fn exp(b: i64, e: i64) -> i64 {
        if e == 0 {
            return 1;
        }

        let half = exp(b, e / 2);
        if e % 2 == 0 {
            return mul(half, half);
        }

        mul(half, mul(half, b))
    }

    /// A trait implementing the unsigned bit shifts.
    pub trait UnsignedShift {
        fn unsigned_shl(self, n: u32) -> Self;
        fn unsigned_shr(self, n: u32) -> Self;
    }

    /// A trait implementing the integer square root.
    pub trait ISqrt {
        fn isqrt(&self) -> Self
        where
            Self: Sized,
        {
            self.isqrt_checked()
                .expect("cannot calculate square root of negative number")
        }

        fn isqrt_checked(&self) -> Option<Self>
        where
            Self: Sized;
    }

    macro_rules! math_traits_impl {
        ($T:ty, $U: ty) => {
            impl UnsignedShift for $T {
                #[inline]
                fn unsigned_shl(self, n: u32) -> Self {
                    ((self as $U) << n) as $T
                }

                #[inline]
                fn unsigned_shr(self, n: u32) -> Self {
                    ((self as $U) >> n) as $T
                }
            }

            impl ISqrt for $T {
                #[inline]
                fn isqrt_checked(&self) -> Option<Self> {
                    use core::cmp::Ordering;
                    match self.cmp(&<$T>::default()) {
                        // Hopefully this will be stripped for unsigned numbers (impossible condition)
                        Ordering::Less => return None,
                        Ordering::Equal => return Some(<$T>::default()),
                        _ => {}
                    }

                    // Compute bit, the largest power of 4 <= n
                    let max_shift: u32 = <$T>::default().leading_zeros() - 1;
                    let shift: u32 = (max_shift - self.leading_zeros()) & !1;
                    let mut bit = <$T>::try_from(1).unwrap().unsigned_shl(shift);

                    // Algorithm based on the implementation in:
                    // https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)
                    // Note that result/bit are logically unsigned (even if T is signed).
                    let mut n = *self;
                    let mut result = <$T>::default();
                    while bit != <$T>::default() {
                        if n >= (result + bit) {
                            n -= result + bit;
                            result = result.unsigned_shr(1) + bit;
                        } else {
                            result = result.unsigned_shr(1);
                        }
                        bit = bit.unsigned_shr(2);
                    }
                    Some(result)
                }
            }
        };
    }

    math_traits_impl!(i8, u8);
    math_traits_impl!(u8, u8);
    math_traits_impl!(i16, u16);
    math_traits_impl!(u16, u16);
    math_traits_impl!(i32, u32);
    math_traits_impl!(u32, u32);
    math_traits_impl!(i64, u64);
    math_traits_impl!(u64, u64);
    math_traits_impl!(i128, u128);
    math_traits_impl!(u128, u128);
    math_traits_impl!(isize, usize);
    math_traits_impl!(usize, usize);
}

#[allow(dead_code)]
mod output {
    pub fn yes() {
        println!("YES");
    }

    pub fn no() {
        println!("NO");
    }

    pub fn yesno(ans: bool) {
        println!("{}", if ans { "YES" } else { "NO" });
    }
}

#[allow(dead_code)]
mod input {
    use std::collections::VecDeque;
    use std::io;
    use std::str::FromStr;

    pub struct Scanner {
        buffer: VecDeque<String>,
    }

    impl Scanner {
        pub fn new() -> Scanner {
            Scanner {
                buffer: VecDeque::new(),
            }
        }

        pub fn next<T: FromStr>(&mut self) -> T {
            if self.buffer.is_empty() {
                let mut input = String::new();

                io::stdin().read_line(&mut input).ok();

                for word in input.split_whitespace() {
                    self.buffer.push_back(word.to_string())
                }
            }

            let front = self.buffer.pop_front().unwrap();
            front.parse::<T>().ok().unwrap()
        }

        pub fn next_vec<T: FromStr>(&mut self, n: usize) -> Vec<T> {
            let mut arr = vec![];

            for _ in 0..n {
                arr.push(self.next::<T>());
            }

            arr
        }
    }
}