2021/05/14

riano_ 師匠とバチャ。

https://kenkoooo.com/atcoder/#/contest/show/9527bbfe-99c2-441c-8afa-43db914ee576

D - Friend Suggestions

Union-Find

A - Robot Racing

まあまあややこしい。その割に最終的にシンプルな実装になる。

E - Second Sum

setを使って区間をあぶり出すやつ。

E - Sorted and Sorted

ひと目DPしたい気持ちになり、転倒数をどうするか、という時に Fenwick's Tree があり助かる。

F - Colorful Tree

 O(N^2) の DP をしたい気持ちになるが、1回の遷移で変化するのがN要素のうち1つだけということに気づくと、オイラーツアーで押し通せる。

F - Square

とにかく大変なDP。

use crate::mod_int::ModInt;
use std::collections::{BTreeMap, BTreeSet};

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());

    // use rand::prelude::*;
    // let mut rng = StdRng::seed_from_u64(114514);
    // loop {
    //     let n = 4;
    //     let mut set = BTreeSet::new();
    //     while set.len() < n {
    //         let i = rng.gen_range(0, n);
    //         let j = rng.gen_range(0, n);
    //         set.insert((i, j));
    //     }
    //
    //     let mut constraints = vec![];
    //     for (i, j) in set {
    //         let v = rng.gen_range(0, 2);
    //         constraints.push((i, j, v));
    //     }
    //
    //     let ans = solve(n, &constraints).value();
    //     let naive = naive(n, &constraints);
    //     if ans != naive {
    //         for (i, j, v) in constraints {
    //             println!("{} {} {}", i, j, v);
    //         }
    //         panic!();
    //     }
    // }
    /*
    4 4
    1 2 1
    3 2 1
    4 1 0
    4 4 1


    */

    let n: usize = sc.read();
    let m: usize = sc.read();
    let mut constraints = Vec::with_capacity(m);
    for _ in 0..m {
        let a = sc.usize0();
        let b = sc.usize0();
        let c: usize = sc.read();
        constraints.push((a, b, c));
    }
    let ans = solve(n, &constraints);
    println!("{}", ans.value());
    // println!("{}", naive(n, &constraints));
}
fn naive(n: usize, constraints: &[(usize, usize, usize)]) -> i64 {
    let mut sum = 0;
    for mask in 0..(1 << (n * n)) {
        let mut mat = vec![vec![0; n]; n];
        for i in 0..(n * n) {
            mat[i / n][i % n] = (mask >> i) & 1;
        }

        let mut ok = true;
        for &(i, j, v) in constraints {
            if mat[i][j] != v {
                ok = false;
            }
        }
        if !ok {
            continue;
        }

        for i in 0..n {
            for j in (i + 2)..=n {
                let sum = mat[i..j]
                    .iter()
                    .flat_map(|row| row[i..j].iter())
                    .sum::<usize>();
                if sum % 2 != 0 {
                    ok = false;
                    break;
                }
            }
            if !ok {
                break;
            }
        }

        if ok {
            sum += 1;
        }
    }
    sum
}

fn solve(n: usize, constraints: &[(usize, usize, usize)]) -> ModInt {
    let mut map = BTreeMap::new();
    let mut other_map = BTreeMap::new();
    let mut filled = vec![0; n];
    for &(a, b, c) in constraints {
        if a.max(b) - a.min(b) <= 2 {
            map.insert((a, b), c);
        } else {
            let pos = a.max(b);
            let t = a.min(b);
            if let Some(&prev) = other_map.get(&(pos, t)) {
                if prev == c {
                    filled[pos] += 1;
                    assert_eq!(other_map.remove(&(pos, t)), Some(c));
                } else {
                    return ModInt::from(0);
                }
            } else {
                other_map.insert((pos, t), c);
            }
        }
    }

    for ((pos, _), _) in other_map {
        filled[pos] += 1;
    }

    let mut pow2 = vec![ModInt::from(1); n + 1];
    for i in 1..=n {
        pow2[i] = pow2[i - 1] * 2;
    }

    let mut dp = vec![ModInt::from(1); 2];
    if let Some(&x) = map.get(&(0, 0)) {
        dp[x ^ 1] = ModInt::from(0);
    }
    for i in 1..n {
        let mut next_dp = vec![ModInt::from(0); 2];
        let (next_from, next_to) = match map.get(&(i, i)) {
            Some(0) => (0, 1),
            Some(1) => (1, 2),
            None => (0, 2),
            _ => unreachable!(),
        };
        let (up_from, up_to) = match map.get(&(i - 1, i)) {
            Some(0) => (0, 1),
            Some(1) => (1, 2),
            None => (0, 2),
            _ => unreachable!(),
        };
        let (left_from, left_to) = match map.get(&(i, i - 1)) {
            Some(0) => (0, 1),
            Some(1) => (1, 2),
            None => (0, 2),
            _ => unreachable!(),
        };
        for prev_cell in 0..2 {
            for next_cell in next_from..next_to {
                for up_cell in up_from..up_to {
                    for left_cell in left_from..left_to {
                        let sum = prev_cell + next_cell + up_cell + left_cell;
                        if sum % 2 != 0 {
                            continue;
                        }

                        if i >= 2 {
                            let c = match (map.get(&(i, i - 2)), map.get(&(i - 2, i))) {
                                (Some(&a), Some(&b)) => {
                                    if (a + b + prev_cell) % 2 == 0 {
                                        1
                                    } else {
                                        0
                                    }
                                }
                                (Some(_), None) | (None, Some(_)) => 1,
                                (None, None) => 2,
                            };
                            assert_eq!(sum & 1, 0);
                            assert!(i - 2 >= filled[i], "{:?} {:?}", constraints, filled);
                            let remain = i - 2 - filled[i];
                            next_dp[next_cell] += dp[prev_cell] * pow2[remain] * c;
                        // eprintln!(
                        //     "{:?} <- {} * {} * {}",
                        //     (i, next_cell),
                        //     dp[prev_cell].value(),
                        //     pow2[remain].value(),
                        //     c
                        // );
                        } else {
                            next_dp[next_cell] += dp[prev_cell];
                        }
                    }
                }
            }
        }
        dp = next_dp;
    }
    dp[0] + dp[1]
}

pub mod mod_int {
    type ModInternalNum = i64;
    thread_local!(
        static MOD: std::cell::RefCell<ModInternalNum> = std::cell::RefCell::new(0);
    );

    pub fn set_mod_int<T: ToInternalNum>(v: T) {
        MOD.with(|x| x.replace(v.to_internal_num()));
    }
    fn modulo() -> ModInternalNum {
        998244353
    }

    pub struct ModInt(ModInternalNum);
    impl Clone for ModInt {
        fn clone(&self) -> Self {
            Self(self.0)
        }
    }
    impl Copy for ModInt {}

    impl ModInt {
        fn internal_new(mut v: ModInternalNum) -> Self {
            let m = modulo();
            if v >= m {
                v %= m;
            }
            Self(v)
        }

        pub fn internal_pow(&self, mut e: ModInternalNum) -> Self {
            let mut result = 1;
            let mut cur = self.0;
            let modulo = modulo();
            while e > 0 {
                if e & 1 == 1 {
                    result *= cur;
                    result %= modulo;
                }
                e >>= 1;
                cur = (cur * cur) % modulo;
            }
            Self(result)
        }

        pub fn pow<T>(&self, e: T) -> Self
        where
            T: ToInternalNum,
        {
            self.internal_pow(e.to_internal_num())
        }

        pub fn value(&self) -> ModInternalNum {
            self.0
        }
    }

    pub trait ToInternalNum {
        fn to_internal_num(&self) -> ModInternalNum;
    }
    impl ToInternalNum for ModInt {
        fn to_internal_num(&self) -> ModInternalNum {
            self.0
        }
    }
    macro_rules! impl_primitive {
        ($primitive:ident) => {
            impl From<$primitive> for ModInt {
                fn from(v: $primitive) -> Self {
                    let v = v as ModInternalNum;
                    Self::internal_new(v)
                }
            }
            impl ToInternalNum for $primitive {
                fn to_internal_num(&self) -> ModInternalNum {
                    *self as ModInternalNum
                }
            }
        };
    }
    impl_primitive!(u8);
    impl_primitive!(u16);
    impl_primitive!(u32);
    impl_primitive!(u64);
    impl_primitive!(usize);
    impl_primitive!(i8);
    impl_primitive!(i16);
    impl_primitive!(i32);
    impl_primitive!(i64);
    impl_primitive!(isize);

    impl<T: ToInternalNum> std::ops::AddAssign<T> for ModInt {
        fn add_assign(&mut self, rhs: T) {
            let mut rhs = rhs.to_internal_num();
            let m = modulo();
            if rhs >= m {
                rhs %= m;
            }

            self.0 += rhs;
            if self.0 >= m {
                self.0 -= m;
            }
        }
    }

    impl<T: ToInternalNum> std::ops::Add<T> for ModInt {
        type Output = ModInt;
        fn add(self, rhs: T) -> Self::Output {
            let mut res = self;
            res += rhs;
            res
        }
    }
    impl<T: ToInternalNum> std::ops::SubAssign<T> for ModInt {
        fn sub_assign(&mut self, rhs: T) {
            let mut rhs = rhs.to_internal_num();
            let m = modulo();
            if rhs >= m {
                rhs %= m;
            }
            if rhs > 0 {
                self.0 += m - rhs;
            }
            if self.0 >= m {
                self.0 -= m;
            }
        }
    }
    impl<T: ToInternalNum> std::ops::Sub<T> for ModInt {
        type Output = Self;
        fn sub(self, rhs: T) -> Self::Output {
            let mut res = self;
            res -= rhs;
            res
        }
    }
    impl<T: ToInternalNum> std::ops::MulAssign<T> for ModInt {
        fn mul_assign(&mut self, rhs: T) {
            let mut rhs = rhs.to_internal_num();
            let m = modulo();
            if rhs >= m {
                rhs %= m;
            }
            self.0 *= rhs;
            self.0 %= m;
        }
    }
    impl<T: ToInternalNum> std::ops::Mul<T> for ModInt {
        type Output = Self;
        fn mul(self, rhs: T) -> Self::Output {
            let mut res = self;
            res *= rhs;
            res
        }
    }

    impl<T: ToInternalNum> std::ops::DivAssign<T> for ModInt {
        fn div_assign(&mut self, rhs: T) {
            let mut rhs = rhs.to_internal_num();
            let m = modulo();
            if rhs >= m {
                rhs %= m;
            }
            let inv = Self(rhs).internal_pow(m - 2);
            self.0 *= inv.value();
            self.0 %= m;
        }
    }

    impl<T: ToInternalNum> std::ops::Div<T> for ModInt {
        type Output = Self;
        fn div(self, rhs: T) -> Self::Output {
            let mut res = self;
            res /= rhs;
            res
        }
    }
}

pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> IO<R, W> {
        IO(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write_all(s.to_string().as_bytes()).unwrap();
    }
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn usize0(&mut self) -> usize {
        self.read::<usize>() - 1
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}