047 - Monochromatic Diagonal(★7) Editorial by ngtkana


公式解説と異なる点は2つです。

  • 公式解説は対角線が \(\mathtt { R }\) ならば \(\mathtt { G }\)\(\mathtt { B }\) を入れ替えるという操作をしているため3パターンすべてで照合してみる必要がありますが、その場合分けがありません。
  • 文字列称号問題を解くにローリングハッシュでも Z-algorithm でもなく、Knuth–Morris–Patt アルゴリズム(以下 KMP アルゴリズム)を使っています。

文字列称号問題への帰着

\(\mathtt { R }, \mathtt { G }, \mathtt { B }\) をそれぞれ \(0, 1, 2\) と表します。マス目 \((i, j)\) の色 \(c\) を固定すると、\(T _ j\)\(S _ i\) を用いて \(T _ j = 2 c - S _ i \bmod 3\) と表せます。従ってマス目 \((i, j), (i + 1, j + 1), \dots ( i + d, j + d)\) が一色のみで塗られている条件は

\[ s _ { i + h + 1 } - s _ { i + h } = - t _ { j + h + 1 } + t _ { j + h } \pmod 3 \quad ( 0 \le j \lt d ) \]

が成り立つことです。

そこで長さ \(N - 1\) の数列 \(A, B\)\(A _ i = S _ { i + 1 } - S _ i, \quad B _ j = - T _ { j + 1 } + T _ j\) で定義すると、\(k\) に対応する対角線のマス目が一色のみで塗られている(※ 以下、「良い」と略記します。)条件は、\(A\)\(B\)\(k\) だけずらして照合したときに空文字でない部分は一致している、すなわち \(1 \le i \le N, 1 + k \le i \le N + k\) なる任意の整数 \(i\) に対して \(S _ i = T _ { i + k}\) の成り立つようなものの個数に一致します。(なおわかるとおり、\(k = -N + 1, N - 1\) のとき、空文字列同士を比較するため必ず \(\mathtt { true }\) になります。)

文字列称号問題の解法

\(0 \le k\) であるものだけ列挙します。\(k \lt 0\) も同様です。

まず \(0 \le k\) な良い対角線であって、\(k\) の最も小さなものを計算しましょう。これは \(A\) をテキスト、\(B\) を検索クエリとする MP アルゴリズムの中で初めて \(A\) の終端に到達した時点の状態を見ればわかります。

次に良い対角線をすべて列挙しましょう。ある良い対角線が与えられたときに、次に長い良い対角線は、今度は(先程は \(A\) について計算しましたが) \(B\)failure functionを計算しておくと、これを見るとわかります。

なお \(B\) の failure function として、いわゆる tagged bordder(MP法とKMP法の違い - 生きたい)を用いると、本来良いはずの対角線をスキップされてしまい困るので、今回は経路圧縮をしない方を使いましょう。(ということは KMP ではなくて MP になるのでしょうか。用語には自信がありません。)

use itertools::Itertools;
use proconio::{fastout, input};
use std::{iter::once, usize::MAX};

#[fastout]
fn main() {
    input! {
        _n: usize,
        s: String,
        t: String,
    }

    // 隣接差の構築
    let a = s
        .chars()
        .map(|c| c as u8 % 3)
        .tuple_windows()
        .map(|(x, y)| (3 + y - x) % 3)
        .collect_vec();
    let b = t
        .chars()
        .map(|c| (240 - c as u8) % 3)
        .tuple_windows()
        .map(|(x, y)| (3 + y - x) % 3)
        .collect_vec();

    // 良い対角線の列挙
    let mut ans = 0;
    let kmp_a = kmp(&a);
    let kmp_b = kmp(&b);
    for (a, b, kmp_a, kmp_b) in once((&a, &b, &kmp_a, &kmp_b)).chain(once((&b, &a, &kmp_b, &kmp_a)))
    {
        // 最も長い良い対角線の探索
        let mut i = 0;
        for j in 0..=a.len() {
            while j != a.len() && a[j] != b[j - i] {
                if i == j {
                    i += 1;
                    break;
                }
                i = j - kmp_a[j - i];
            }
        }
        // 良い対角線の列挙
        let mut l = a.len() - i;
        while l != MAX {
            l = kmp_b[l];
            ans += 1;
        }
    }
    if a == b {
        ans -= 1;
    }
    println!("{}", ans);
}

fn kmp(s: &[u8]) -> Vec<usize> {
    let mut a = vec![!0; s.len() + 1];
    let mut j = MAX;
    for (i, c) in s.iter().enumerate() {
        while j != MAX && &s[j] != c {
            j = a[j];
        }
        j = j.wrapping_add(1);
        a[i + 1] = j;
    }
    a
}

提出 (56 ms): https://atcoder.jp/contests/typical90/submissions/28326658

posted:
last update: