Codeforces Round #344 (Div. 2): D. Messenger

https://codeforces.com/problemset/problem/631/D

問題概要

文字列S, Tが与えられる。SがTを部分文字列として何個含むか求めよ。ただしS, Tは(回数, 文字)というタプルの列として与えられる。これはその文字が何回連続するかを表した形式で、例えば(3, a)(2, b)(4, c)はaaabbccccという文字列を表す。

S, Tに対応するタプル列の長さ <= 2×105

各タプル中の「回数」の値 <= 106

解法

まず表記の形式をユニークにするため、同じ文字に対応するタプルが連続しているところはひとつにまとめてしまう((3, a)(4, a)(1, a)のような部分は(8, a)にまとめる)。このようにすると、TがSの部分文字列として現れるならば対応箇所のタプル列の長さがSとTで等しくなる。これによって以下のような場合分けができる。

Tのタプル列の長さが1の場合

これは単純にSの要素をひとつずつ見ていって愚直に数えれば良い。

Tのタプル列の長さが2の場合

これもSの要素を2つずつ見ていけば数えられる。

Tのタプル列の長さが3以上の場合

このときSのある箇所がTの部分文字列であるための条件は

  • すべてのタプルについて、文字が一致する
  • 両端以外のタプルについて、回数が一致する
  • 両端のタプルについて、Sの回数がTの回数以上である

のすべてが満たされることである。そしてこのような部分をSから見つけ出すためには

  • 文字だけに着目して文字列検索を行い、Sの中でTと一致する箇所を列挙する
  • (両端以外の)回数だけに注目して「数字列」の検索を行い、Sの中でTと一致する箇所を列挙する

以上2つの前処理を行っておけばよい。前者で列挙した文字列一致インデックスを舐めていき、後者で列挙した数字列インデックスに対応する箇所(両端以外の列)が存在しているか、またS側の両端の回数はTの両端の回数と一致しているかを確認する。すべて満たされていればカウント+1となる。

肝心のこの前処理をどう効率的に行うかだが、これは単純な文字列検索なのでKMP法によってO(len(S) + len(T))でできる。

感想

KMP法を初めて書いたので記念カキコ(wikipedia擬似コードを丸写ししただけ)

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, core.stdc.stdio;

void main() {
    auto s = readln.split.map!(to!int);
    long[] A_INT, B_INT;
    string A_STR, B_STR;

    void parse(ref long[] x, ref string y) {
        auto t = readln.split;
        foreach (a; t) {
            auto v = a.split("-").front.to!long;
            auto c = a.split("-").back[0];
            if (y.empty || y.back != c) {
                y ~= c;
                x ~= v;
            } else {
                x.back += v;
            }
        }
    }

    parse(A_INT, A_STR);
    parse(B_INT, B_STR);
    int N = A_INT.length.to!int;
    int M = B_INT.length.to!int;
    long ans = 0;

    if (M == 1) {
        foreach (i; 0..N) {
            if (A_STR[i] != B_STR.front) continue;
            ans += max(0L, A_INT[i] - B_INT.front + 1);
        }
    } else if (M == 2) {
        foreach (i; 0..N-1) {
            if (A_STR[i..i+2] != B_STR) continue;
            if (A_INT[i] >= B_INT[0] && A_INT[i+1] >= B_INT[1]) ans += 1;
        }
    } else {
        auto int_idx = kmp_search(A_INT, B_INT[1..M-1]);
        auto str_idx = kmp_search(A_STR, B_STR);
        bool[int] dict;
        foreach (idx; int_idx) dict[idx-1] = true;
        foreach (idx; str_idx) {
            if (A_INT[idx] >= B_INT[0] && A_INT[idx+M-1] >= B_INT[M-1] && idx in dict) {
                ans += 1;
            }
        }
    }

    ans.writeln;
}

int[] kmp_search(T)(const T[] S, const T[] W) {
    int j, k;
    int[] P;
    auto table = kmp_table(W);

    while (j < S.length) {
        if (W[k] == S[j]) {
            j += 1;
            k += 1;
            if (k == W.length) {
                P ~= j - k;
                k = table[k];
            }
        } else {
            k = table[k];
            if (k < 0) {
                j += 1;
                k += 1;
            }
        }
    }

    return P;
}

int[] kmp_table(T)(const T[] W) {
    auto n = W.length.to!int;
    auto table = new int[](n+1);
    table[0] = -1;
    int pos = 1, cnd = 0;

    while (pos < n) {
        if (W[pos] == W[cnd]) {
            table[pos] = table[cnd];
        } else {
            table[pos] = cnd;
            cnd = table[cnd];
            while (cnd >= 0 && W[pos] != W[cnd]) {
                cnd = table[cnd];
            }
        }
        pos += 1;
        cnd += 1;
    }

    table[pos] = cnd;
    return table;
}