#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>;
   |      ^^^^^

ソースコード

raw source code

// 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 ----------

yukicoder

結果

コンパイルメッセージ

ソースコード