yukicoder No.770 Median Sequence

思いついた解法が嘘解法で、正しい解法を思いつくためにかなり多量の実験をした。
No.770 Median Sequence - yukicoder

問題

正の整数Nが与えられる。要素が1からNまでの長さNの整数列bに対して、以下の操作で数列aを作る。

a[1] = med(1, b[1], N)
a[2] = med(med(1, b[1], N - 1), b[2], N)
a[3] = med(med(med(1, b[1], N - 2), b[2], N - 1), b[3], N)
...
a[i] = (m = 1; for(int j = 1; j <= i; ++j) m = med(m, b[i], N - i + j); を実行したときのmの値)

このとき、aとしてあり得るものの個数を数えよ。

解法

aとしてあり得るものは以下の操作で作れるものであり、およびそれに限る。
ここで、数列{c[1], ..., c[m]} が減少的であるとは、最後にc[i] < c[i+1]が成立する添字より最後にc[i] > c[i+1]が成立する添字が大きいことをいうことにする。
どちらかが存在しないときは、存在する方が大きいものとする。どちらも存在しないときはどちらも大きくないものとする。

a[1] は任意
数列{a[1], ..., a[i - 1]}が減少的である場合、あるいはk = N - a[i - 1]としてi > k, a[i - k] = a[i - k + 1] = ... = a[i - 1]の両方が成り立つ場合、a[i] >= a[i - 1] - 1となるようにa[i]を選ぶ
それ以外の場合、a[i] >= a[i - 1]となるようにa[i]を選ぶ

これは以下のようなDPで計算できる。

DP[i][j]: 長さiで最後の要素がjであるような、減少的でない数列の個数
DP2[i][j]: 長さiで最後の要素がjであるような、減少的である数列の個数

実装上の注意点

遷移が合っているかどうかを確認する
答えが合わなかったら愚直解法と比較する

提出: #305095 No.770 Median Sequence - yukicoder (Rust)
#305093 No.770 Median Sequence - yukicoder (愚直解法のコードを含む)

const MOD: i64 = 1_000_000_007;
define_mod!(P, MOD);
type ModInt = mod_int::ModInt<P>;

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($format:expr) => (write!(out,$format).unwrap());
        ($format:expr, $($args:expr),+) => (write!(out,$format,$($args),*).unwrap())
    }
    input! {
        n: usize,
    }
    let mut dp = vec![vec![ModInt::new(0); n]; n + 1];
    let mut dp2 = vec![vec![ModInt::new(0); n]; n + 1];
    dp[0][0] = ModInt::new(1);
    for i in 1 .. n + 1 {
        let mut cur = ModInt::new(0);
        for j in 0 .. n {
            cur += dp[i - 1][j];
            dp[i][j] += cur;
            cur += dp2[i - 1][j];
        }
        for j in 0 .. n - 1 {
            let turn = n - 1 - j;
            if i >= turn {
                dp2[i][j] += dp[i - turn][j + 1]
                    + dp2[i - 1][j + 1] + dp2[i - 1][j];
            }
        }
    }
    let mut tot = ModInt::new(0);
    for i in 0 .. n { tot += dp[n][i] + dp2[n][i]; }
    puts!("{}\n", tot);
}

まとめ

最初a[i]>=a[i-1]-1ならOKだと思ってサンプルが盛大に合わず苦しんだ。(そのような数列の個数はN=6のとき4488個ある。)