2021/05/14
riano_ 師匠とバチャ。
https://kenkoooo.com/atcoder/#/contest/show/9527bbfe-99c2-441c-8afa-43db914ee576
Union-Find
まあまあややこしい。その割に最終的にシンプルな実装になる。
setを使って区間をあぶり出すやつ。
ひと目DPしたい気持ちになり、転倒数をどうするか、という時に Fenwick's Tree があり助かる。
の DP をしたい気持ちになるが、1回の遷移で変化するのがN要素のうち1つだけということに気づくと、オイラーツアーで押し通せる。
とにかく大変な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() } }