AtCoder Regular Contest 071 F - Infinite Sequence

AtCoder 競技プログラミング 動的計画法 数学

問題

rated 黄色 diff 最後の生き残り

F - Infinite Sequence

解法

1より大きい数が連続すると以降は全て決まるので、

  • (1)(211)(3111)(41111)... から構成される prefix

と

  • a>1 かつ b>1 として abbbbbb... となる suffix
  • a>1として a11111.... となる suffix

のいずれかの suffix を組み合わせて作る。

dp[i] := 先頭の i 要素を (1)(211)(3111)... から構成する組み合わせの数とすると dp[i] = dp[i-1] + dp[i-3] + dp[i-4] + ... + dp[0] であることが分かるので、dp[i] が求まる。

(1)(211)(3111)(41111)... から構成される prefix と a>1 かつ b>1 として abbbbbb... となる suffix の組み合わせの数は

    for i in 0..(n - 1) {
        ans += dp[i] * (n - 1) * (n - 1);
    }

で求めている。

また、n-1項を(1)(211)(3111)(41111)... から構成し、n項目を好きに選んで無限に繰り返すのはdp[n - 1] * nとなる。これを (1) とする。

最初の k 要素を(1)(211)(3111)(41111)... から構成し、その次の要素を a>1 とし、それ以降を 1111... とするような数列を求めたいが、(1)で選んだものは取り除きたい。その場合、先頭k要素と次のaとそれに続くa個の1を合わせた k+1+a 要素がn個以上になっていなければならない（n-1個以下の場合、dp[k+1+a]に含まれているため）し、kかつk<=n-2でなければならない（k=n-1はdp[n-1]*nに含まれているため）。

よって k+1+a>=n かつ k<=n-2かつa>1なので、n >= a > max(n-k-2,1) より、kを決めた時に a として選べるものはmin(k+2, n-1)個である。よってこの場合の組み合わせは

    for k in 0..(n - 1) {
        ans += dp[k] * min(k + 2, n - 1);
    }

で求まる。

コード

use self::mod_int::ModInt;
use std::cmp::min;

const MOD: usize = 1e9 as usize + 7;

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());
    let n: usize = sc.read();
    if n == 1 {
        println!("1");
        return;
    }

    let mut dp = vec![ModInt(0); n];
    dp[0] = ModInt(1);
    dp[1] = ModInt(1);

    let mut sum = vec![ModInt(0); n + 1];
    sum[1] = dp[0];
    sum[2] = dp[1] + sum[1];
    for i in 2..n {
        dp[i] = dp[i - 1] + sum[i - 2];
        sum[i + 1] = sum[i] + dp[i];
    }

    let mut ans = ModInt(0);
    for i in 0..(n - 1) {
        ans += dp[i] * (n - 1) * (n - 1);
    }
    ans += dp[n - 1] * n;
    for k in 0..(n - 1) {
        ans += dp[k] * min(k + 2, n - 1);
    }
    println!("{}", ans.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)]
    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, mut rhs: Num) -> ModInt<Num> {
            if rhs >= MOD {
                rhs %= MOD;
            }
            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, mut rhs: Num) -> ModInt<Num> {
            if rhs >= MOD {
                rhs %= MOD;
            }
            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, mut rhs: Num) -> ModInt<Num> {
            if rhs >= MOD {
                rhs %= MOD;
            }
            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 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) -> Self {
        IO(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write(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 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()
    }
}