DISCO presents ディスカバリーチャンネル コードコンテスト2019 本戦 D - DISCO!

まあ実家だよね。
D - DISCO!

問題

'D', 'I', 'S', 'C', 'O'のみからなる文字列Sが与えられる。このとき、以下のQ個の質問に答えよ:

  • SのL_i文字目からR_i文字目までからなる部分文字列を考える。この文字列に部分列として含まれる"DISCO"は何通りあるか? つまり、単調増加な添字の5つ組であって、それぞれの位置の文字を並べると"DISCO"になるようなものは何通りあるか? 答えを2^32で割った余りを求めよ。

解法

まずは「Sの"DISCO"という部分文字列を数え上げる問題」を考えよう。これは以下のような6 * |S|エントリのDPでできる:

  • DP[i][j]: i文字目まで見たときに、"DISCO"の何文字目までできているか (0 <= i <= |S|, 0 <= j <= 5)
  • 遷移: i文字目を取るか取らないかの2通り。ただしi文字目を取れるのはS[i] == "DISCO"[j]のときだけ。

こうすると、例えばS[i] = 'D'の時の遷移は

(DP[i][0] ... DP[i][5]) =
(DP[i - 1][0] ... DP[i - 1][5])
(1 0 0 0 0 0
1 1 0 0 0 0
0 0 1 0 0 0
0 0 0 1 0 0
0 0 0 0 1 0
0 0 0 0 0 1)

と、行ベクトルに右から行列を掛ける形で表せる。最終的に欲しい値は、(DP[0][0] ... DP[0][5]) = (1 0 0 0 0 0)としたときのDP[|S|][5]である。
よって、文字S[i]に対応する行列をA[i]と置くと、(1 0 0 0 0 0)A[1]A[2] ... A[|S|]の第5成分が欲しい値である。
今は1文字目から|S|文字目までの遷移を考えたが、欲しいのは行列積なので、任意の区間[L, R]に対しA[L] A[L + 1] ... A[R] (の0行5列成分)が分かればよい。これは行列をセグメント木に載せればできる。計算量は愚直にやるとO(K^3 (|S| + Q) log |S|)であり、制限時間が長いのでこれで十分である。

登場する典型

  • DPの遷移が線形変換であるときに、高速化をする定跡
    • 行列累乗、複数の変換を行列積に置き換えることでできる高速化(セグメント木に載せるなど)

実装上の注意点

  • 制限時間が13秒と長いので、間に合う中で一番簡単な実装を選択する。今回の場合、登場する行列が全て可逆なので、行列積バージョンの累積和を取るO(K^2(|S| + Q))の解法を書くこともできるが、セグメント木で殴るO(K^3 (|S| + Q) log |S|)の解法でもよい。
    • なお、遷移を両側から行うために、"DISCO"の部分文字列10通りについて遷移を手書きする方法があるが、より簡単な実装を考えるべきである。
  • Rustでオーバーフローする演算を行いたいときは、core::num::Wrappingを使うか基本型のwrapping_***メソッドを使う。
    • 間違えてi32で演算しないようにする(1敗)

提出: Submission #4039621 - DISCO presents ディスカバリーチャンネル コードコンテスト2019 本戦 (Rust)

// SegTree省略

extern crate core;
use core::num::Wrapping;
const SZ: usize = 6;
type Mat = [[Wrapping<u32>; 6]; 6];
 
fn mul(a: Mat, b: Mat) -> Mat {
    let mut ans = [[Wrapping(0); 6]; 6];
    for i in 0 .. SZ {
        for j in 0 .. SZ {
            for k in 0 .. SZ {
                ans[i][k] += a[i][j] * b[j][k];
            }
        }
    }
    ans
}
 
fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        s: chars,
        lr: [(usize1, usize)],
    }
    let n = s.len();
    let mut unit: Mat = [[Wrapping(0); 6]; 6];
    for i in 0 .. SZ { unit[i][i] = Wrapping(1); }
    let mut st = SegTree::new(n, |x, y| mul(x, y), unit);
    for i in 0 .. s.len() {
        let idx = match s[i] {
            'D' => 0,
            'I' => 1,
            'S' => 2,
            'C' => 3,
            'O' => 4,
            _ => unreachable!(),
        };
        let mut m = unit;
        m[idx][idx + 1] = Wrapping(1);
        st.update(i, m);
    }
    for (l, r) in lr {
        let ans = st.query(l, r);
        puts!("{}\n", ans[0][5]);
    }
}

まとめ

この手の解ける問題で方針をミスらなかったり実装を単純化したりするのは極めて大事。