#1066973 (Rust) No.1502 Many Simple Additions

• 問題: No.1502 Many Simple Additions
• ユーザー: akakimidori
• 提出日時: 2025-04-13 14:54:44
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 81 ms / 2,000 ms
• コード長: 6,615 bytes
• コンパイル時間: 14,362 ms
• コンパイル使用メモリ: 397,608 KB
• 実行使用メモリ: 18,216 KB
• 最終ジャッジ日時: 2025-04-13 14:55:02
• 合計ジャッジ時間: 17,332 ms
• ジャッジサーバーID （参考情報）: judge1 / judge5

ファイルパターン	結果
sample	AC * 5
other	AC * 39

コンパイルメッセージ

warning: unused import: `std::io::Write`
  --> src/main.rs:10:5
   |
10 | use std::io::Write;
   |     ^^^^^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` on by default

warning: type alias `Map` is never used
  --> src/main.rs:13:6
   |
13 | type Map<K, V> = BTreeMap<K, V>;
   |      ^^^
   |
   = note: `#[warn(dead_code)]` on by default

warning: type alias `Set` is never used
  --> src/main.rs:14:6
   |
14 | type Set<T> = BTreeSet<T>;
   |      ^^^

warning: type alias `Deque` is never used
  --> src/main.rs:15:6
   |
15 | type Deque<T> = VecDeque<T>;
   |      ^^^^^

ソースコード

// A_x + A_y = z
// なる条件M個
// 1 <= A_i <= K
// となるようなものの個数がわかればいい
// まず矛盾しないかというのを解くがこれは関係式を管理すればいい
// 各連結成分について、一つ値を決めると残りも決まる
// 
// 1 <= A_1 <= K のうち何個が条件満たすのか？

use std::io::Write;
use std::collections::*;

type Map<K, V> = BTreeMap<K, V>;
type Set<T> = BTreeSet<T>;
type Deque<T> = VecDeque<T>;

fn main() {
    input! {
        n: usize,
        m: usize,
        k: i64,
        e: [(usize1, usize1, i64); m],
    }
    const MOD: i64 = 1_000_000_007;
    let calc = |k: i64| -> i64 {
        let mut dsu = DSU::new(2 * n);
        for &(x, y, z) in e.iter() {
            if let Some(v) = dsu.find(n + y, x) {
                if v != z {
                    return 0;
                }
            }
            dsu.add_edge(n + y, x, z);
            dsu.add_edge(y, n + x, -z);
        }
        let mut list = vec![vec![]; 2 * n];
        for i in 0..(2 * n) {
            let root = dsu.root(i);
            list[root].push(i);
        }
        let mut ans = 1;
        for i in 0..n {
            let r = dsu.root(i);
            if list[r].is_empty() {
                continue;
            }
            let cond = std::mem::take(&mut list[r]);
            let mut low = 1;
            let mut up = k;
            for v in cond {
                let d = dsu.find(i, v).unwrap();
                if v == n + i {
                    if d > 0 || d % 2 != 0 {
                        return 0;
                    }
                    let v = -d / 2;
                    low.chmax(v);
                    up.chmin(v);
                } else if v < n {
                    low.chmax(1 - d);
                    up.chmin(k - d);
                } else {
                    low.chmax(-k - d);
                    up.chmin(-1 - d);
                }
            }
            if low > up {
                return 0;
            }
            ans = ans * (up - low + 1) % MOD;
            let r = dsu.root(i + n);
            list[r].clear();
        }
        ans
    };
    let ans = (calc(k) - calc(k - 1) + MOD) % MOD;
    println!("{}", ans);
}

impl DSUGroup for i64 {
    fn merge(&self, rhs: &Self) -> Self {
        *self + *rhs
    }
    fn inv(&self) -> Self {
        -*self
    }
    fn e() -> Self {
        0
    }
}

// ---------- begin input macro ----------
// reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
#[macro_export]
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

#[macro_export]
macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};
    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

#[macro_export]
macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };
    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };
    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };
    ($iter:expr, bytes) => {
        read_value!($iter, String).bytes().collect::<Vec<u8>>()
    };
    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };
    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}
// ---------- end input macro ----------
// 関係の木を管理する
// add_edge(src, dst, v)
//   src, dst が連結でない時
//   src->dst 方向に読むとvとなるような辺を追加する
//   dst->src だとv^(-1)
// find(src, dst):
//   src, dst が連結名時
//   srcからdstに向けて関係の木を辿って行った時の積を計算
// 
// verify:
// https://judge.yosupo.jp/submission/259683
// https://judge.yosupo.jp/submission/259688
// ---------- begin DSU Group ---------
pub trait DSUGroup: Clone + Eq {
    fn merge(&self, rhs: &Self) -> Self;
    fn inv(&self) -> Self;
    fn e() -> Self;
}

pub struct DSU<T> {
    p: Vec<i32>,
    data: Vec<T>,
}

impl<T> DSU<T>
where
    T: DSUGroup,
{
    pub fn new(size: usize) -> Self {
        Self {
            p: vec![-1; size],
            data: vec![T::e(); size],
        }
    }
    pub fn add_edge(
        &mut self,
        mut src: usize,
        mut dst: usize,
        mut val: T,
    ) -> Option<(usize, usize)> {
        assert!(src < self.p.len() && dst < self.p.len());
        let p = self.root(src);
        let q = self.root(dst);
        if p == q {
            return None;
        }
        if self.p[p] < self.p[q] {
            std::mem::swap(&mut src, &mut dst);
            val = val.inv();
        }
        let (u, a) = self.root_with(src);
        let (v, b) = self.root_with(dst);
        self.data[u] = a.inv().merge(&val).merge(&b);
        self.p[v] += self.p[u];
        self.p[u] = v as i32;
        Some((v, u))
    }
    pub fn find(&self, src: usize, dst: usize) -> Option<T> {
        assert!(src < self.p.len() && dst < self.p.len());
        if self.root(src) != self.root(dst) {
            return None;
        }
        let (_, p) = self.root_with(src);
        let (_, q) = self.root_with(dst);
        Some(p.merge(&q.inv()))
    }
    fn root(&self, mut v: usize) -> usize {
        while self.p[v] >= 0 {
            v = self.p[v] as usize;
        }
        v
    }
    fn root_with(&self, mut v: usize) -> (usize, T) {
        let mut a = T::e();
        while self.p[v] >= 0 {
            let data = &self.data[v];
            a = a.merge(&data);
            v = self.p[v] as usize;
        }
        (v, a)
    }
}
// ---------- end DSU Group ---------
// ---------- begin chmin, chmax ----------
pub trait ChangeMinMax {
    fn chmin(&mut self, x: Self) -> bool;
    fn chmax(&mut self, x: Self) -> bool;
}

impl<T: PartialOrd> ChangeMinMax for T {
    fn chmin(&mut self, x: Self) -> bool {
        *self > x && {
            *self = x;
            true
        }
    }
    fn chmax(&mut self, x: Self) -> bool {
        *self < x && {
            *self = x;
            true
        }
    }
}
// ---------- end chmin, chmax ----------