AGC039 C - Division by Two with Something - 宇宙ツイッタラーXの憂鬱

問題

解法

各操作を観察すると、どちらも「最も下のビットをpopして反転して最も上にpushする」と言い換えられる。これをN回繰り返すと元の数のビットを反転させたものが得られるので、2N回繰り返すと必ず元の数が得られる。よって繰り返して元の数が得られる繰り返し回数kは2Nの約数であると考えられる。

繰り返し回数をkとし、繰り返されている長さkのビット列を[..., c, b, a] とすると、k回操作した後のビット列は[..., !c, !b, !a, ..., c, b, a, ..., c, b, a, ..., ..., c, b, a] という感じになっている。ここでNがkの倍数だとするとN回繰り返した後に元の数が得られるはずだが、N回操作後のビット列は[..., c!, b!, !a]となっていて元に戻ってはいないため、矛盾する。よってkはNの約数ではない。kは2Nの約数ではあるため、kは偶数であり、2N/kは奇数である。

2N/kが奇数であることから長さkのビット列は[..., !c, !b, !a, ..., c, b, a]となっているはずである。よってこの長さkのビット列は長さk/2のビット列と、それを反転したものを並べたものだと分かる。

各kについて動的計画法で数え上げを計算し、重複したものを包除原理で取り除いてやることで、答えが求まる。

コード

use mod_int::ModInt;
use std::collections::BTreeMap;

const MOD: usize = 998244353;
fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let x = sc
        .chars()
        .into_iter()
        .map(|c| if c == '0' { 0 } else { 1 })
        .collect::<Vec<_>>();

    let m = 2 * n;
    let mut ks = vec![];
    for k in 2..(m + 1) {
        if m % k == 0 && k % 2 == 0 && (m / k) % 2 == 1 {
            ks.push(k);
        }
    }

    let mut ans = vec![ModInt(0); m + 1];
    for &k in ks.iter() {
        let mut dp = vec![vec![ModInt(0); 2]; 2];
        dp[0][0] = ModInt(1);
        for i in 0..(k / 2) {
            let mut next_dp = vec![vec![ModInt(0); 2]; 2];
            for prev in 0..2 {
                for next in 0..2 {
                    for smaller in 0..2 {
                        if smaller == 0 && next > x[i] {
                            continue;
                        }
                        let next_smaller = if next < x[i] { 1 } else { smaller };
                        next_dp[next][next_smaller] += dp[prev][smaller];
                    }
                }
            }
            dp = next_dp;
        }

        let same = dp[0][0] + dp[1][0];
        let smaller = dp[0][1] + dp[1][1];
        let mut ok = true;
        for i in 0..n {
            let s = i % k;
            let t = i % (k / 2);
            let y = if s == t {
                x[i % (k / 2)]
            } else {
                x[i % (k / 2)] ^ 1
            };
            if x[i] > y {
                break;
            }
            if x[i] < y {
                ok = false;
                break;
            }
        }
        ans[k] = if ok { same + smaller } else { smaller };
    }

    let mut result = ModInt(0);
    for k in 1..(m + 1) {
        let mut cur = k * 2;
        while cur <= m {
            if ans[cur].0 > 0 {
                ans[cur] -= ans[k];
            }
            cur += k;
        }
        result += ans[k] * k;
    }
    println!("{}", result.0);
}

pub mod mod_int {
    use super::MOD;
    use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};

    type Num = usize;

    #[derive(Clone, Copy, Debug)]
    pub struct ModInt<T: Copy + Clone>(pub T);

    impl Add<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self + rhs.0
        }
    }

    impl Add<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn add(self, rhs: Num) -> ModInt<Num> {
            let mut t = rhs + self.0;
            if t >= MOD {
                t = t - MOD;
            }
            ModInt(t)
        }
    }

    impl Sub<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: Num) -> ModInt<Num> {
            let rhs = if rhs >= MOD { rhs % MOD } else { rhs };
            let value = if self.0 < rhs { self.0 + MOD } else { self.0 };
            ModInt(value - rhs)
        }
    }

    impl Sub<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn sub(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self - rhs.0
        }
    }

    impl AddAssign<Num> for ModInt<Num> {
        fn add_assign(&mut self, other: Num) {
            *self = *self + other;
        }
    }
    impl AddAssign<ModInt<Num>> for ModInt<Num> {
        fn add_assign(&mut self, other: ModInt<Num>) {
            *self = *self + other;
        }
    }

    impl SubAssign<Num> for ModInt<Num> {
        fn sub_assign(&mut self, other: Num) {
            *self = *self - other;
        }
    }

    impl SubAssign<ModInt<Num>> for ModInt<Num> {
        fn sub_assign(&mut self, other: ModInt<Num>) {
            *self = *self - other;
        }
    }

    impl Div<Num> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: Num) -> ModInt<Num> {
            self * ModInt(rhs).pow(MOD - 2)
        }
    }

    impl Div<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;
        fn div(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self / rhs.0
        }
    }

    impl DivAssign<Num> for ModInt<Num> {
        fn div_assign(&mut self, rhs: Num) {
            *self = *self / rhs
        }
    }
    impl DivAssign<ModInt<Num>> for ModInt<Num> {
        fn div_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self / rhs
        }
    }

    impl Mul<ModInt<Num>> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: ModInt<Num>) -> ModInt<Num> {
            self * rhs.0
        }
    }
    impl Mul<Num> for ModInt<Num> {
        type Output = ModInt<Num>;

        fn mul(self, rhs: Num) -> ModInt<Num> {
            let t = (self.0 * rhs) % MOD;
            ModInt(t)
        }
    }

    impl MulAssign<Num> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: Num) {
            *self = *self * rhs;
        }
    }

    impl MulAssign<ModInt<Num>> for ModInt<Num> {
        fn mul_assign(&mut self, rhs: ModInt<Num>) {
            *self = *self * rhs;
        }
    }

    impl ModInt<Num> {
        pub fn pow(self, e: usize) -> ModInt<Num> {
            let mut result = ModInt(1);
            let mut cur = self;
            let mut e = e;
            while e > 0 {
                if e & 1 == 1 {
                    result *= cur;
                }
                e >>= 1;
                cur *= cur;
            }
            result
        }
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    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()
    }
}