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; }