ARC055C ABCAC - Zアルゴリズム

シェアする

  • このエントリーをはてなブックマークに追加

こんにちは。競プロシリーズ3本目ですね。

今回は、ARC055のC問題「ABCAC」を解きました。

文字列の問題です。

長さ200,000以下の文字列Sが与えられ、「Sを5つの(空でない)部分文字列に分割したとき、1つ目と4つ目、2つ目と5つ目が同じ文字列になるような分割のしかた」の個数を計算する問題です。

例えば、"takaitai"であれば、ta, ka, i, ta, i のように分割する方法と、t, ak, ai, t, ai のように分割する方法の2通りがあります。

文字列A,B,Cを用いて、Sを"ABCAC"として構成できるようなA,B,Cの数とも言えますね。

この問題では「Zアルゴリズム」と呼ばれる文字列の接頭辞に関するアルゴリズムが使えます。

では、考えていきましょう。

愚直に全通りは試せない

Sの頭から何文字かをA、その後の何文字かをB、その後の何文字かをCとして、残りの部分がACと一致するかを調べるという手段は、さすがに間に合いません。

"takaitai"の例を見てみると、どちらの分け方も"takai"と"tai"に分けられています。

この分け方は"ABC"と"AC"に分けていることになります。

文字列Sを二つに分割して、前半をX、後半をYとする(つまりS=XY)と、

・XとYの前からa文字(aは1以上)が一致(この部分がA)

・XとYの後ろからc文字(cは1以上)が一致(この部分がC)

・a+cはYの長さと一致

となれば、Xの先頭a文字と後ろc文字を除いた文字列をBとして(Bの長さが1以上である必要あり)、X=ABC、Y=ACより、S=ABCACと表現できます。

つまり、Sを二つに分ける分け方(Sの長さ程度の数だけある。これをABC/ACの切れ目とする)のすべてについて、上記の条件を満たすかどうかを高速に調べれば解けることになります。

共通接頭辞アルゴリズム Z-algorithm

Zアルゴリズムでは、int Z[]という配列を用意します(実装はvector<int>)。

Z[i]には、文字列Sに対して「SとS[i]以降の部分の共通接頭辞の長さ」が格納されます。

例えば、

S=aaabcdabaacaaabd
Z=*210001021042100

というようになります。

S=aaab...に対して、S[1]以降はaab...となるので、先頭2文字が一致しているのでZ[i]=2。

明らかに、S[i]がS[0]と一致しない場合はZ[i]=0です。

また、Z[0]は、定義通りに書けばSの長さと一致するはず(S[0]以降の文字列はSそのものなので当然全文字一致)ですが、実装の都合上、部分列として扱いたくない場合は0とか-1とかにすることも考えられます。だいたいの場合でZ[0]の値は何でもいいようです。

ZアルゴリズムではこのZ[]を高速に計算することができます。

上記の例で、Sの先頭4文字が"aaab"なので、Zの先頭4つは"*210"となっています。

ここで、Z[11]=4に着目します。

S="aaabc...aaabd"となっており、S[11]~S[14]が"aaab"となり、S[0]~S[3]と一致しており、さらにS[4]とS[15]が異なるので、Z[11]=4となります。

では、Z[12]はどうなるでしょうか。

Z[11]=4より、Z[11]からの4文字が先頭4文字と一致しており、5文字目が異なっているので、

S[11], S[12], S[13], S[14] = S[0], S[1], S[2], S[3]

が成り立ち、かつ、S[15] ≠ S[4] が成り立つことが分かります。

ここで、Z[1]=2に着目します。

Z[1]が2なのは、「S[1]=S[0]、S[2]=S[1]、S[3]≠S[2]」だからです。

ここで、Z[11]=4から得られる情報を利用すると、

S[12]=S[1]=S[0]、S[13]=S[2]=S[1]、S[14]=S[3]≠S[2]

であることが分かります。すると、Z[12]=Z[1]=2であることが分かります。

このように、「先頭と複数文字一致している箇所について、その内部のZ[]の値は、先頭の方にあるZ[]の値を再利用して得られる」ということが分かります。

この事実を基に考えると、

・S[i]≠S[0]ならばZ[i]=0であり、このとき比較は1回で済むのでO(1)

・上記以外の場合は、一致文字数が多いほど比較が必要になるが、一致文字数が多いほど、それ以降の多くのZ[]がO(1)で計算できる

ということが分かるので、ZアルゴリズムによるZ[]の計算量はだいたいO(|S|)程度であることが分かります。

今回の問題では、事前にこのZを計算しておけば、答えの計算もO(|S|)でできるようになります。

今回の問題での考え方

さて、今回の問題の解説の大枠は公式の解説PDFを読んでいただくとして、必要な計算は

XとYの前方一致文字数(これがa)、後方一致文字数(これがc)

を求めることです。

これさえできれば解けます。

さて、文字列Sを S[i] と S[i+1] の間で分割することを考えると、XとYの前方一致文字数は Z[i+1] で得られます。Xの先頭はSの先頭、Yの先頭はS[i+1]だからです。

また、後方一致文字数については、Sを反転させた文字列Tを用意しておくと、Tに関するZの Z[L-i-1] で得られます(L=|S|とおいている)。Tにおける分割箇所が T[L-i-2] と T[L-i-1] になるからです。

ただし、一致部分がY全体になると問題が発生します。

例えば、S="abcdc abc" を X="abcdc" と Y="abc" に分けたとき、XとYの前方一致文字数が3になりますが、これをAの文字数とすると、A="abc"、C="" となってしまい、Cが空文字列になってしまうからです。

したがって、前方・後方一致文字数はYの文字数(|Y|=L-i-1)以下でないといけません。

a=min(L-i-1, Zs[i+1]), c=min(L-i-1, Zt[L-i-1]) となります。

これを基に、S[i] と S[i+1] の間で分割したときのABCの表し方の数を

max(0, a + b - (L - i - 1) + 1)

で考えることができます(0とのmaxを取っているのは、不可能な分割箇所ではこの数字がマイナスになるので、その影響を消すため)。

また、ABC/ACの分割箇所について、Bの長さが1以上なので、必ず|X|>|Y|となります。イコールもあり得ません。

したがって、分割箇所の i は L/2 以上である必要があります。

以上を実装したのが以下のC++コードになります。

#include <iostream>
#include <cstdio>
#include <string>
#include <algorithm>
#include <vector>
 
#define repl(i,a,b) for(int i=(int)(a);i<(int)(b);i++)
#define rep(i,n) repl(i,0,n)
 
typedef long long ll;
 
using namespace std;
 
void Z_algorithm(string str, vector<int> &Z) {
	const int L = str.length();
	for (int i = 1, left = 0, right = 0; i < L; i++) { if (i > right) {
			left = right = i;
			for (; right < L && str[right - left] == str[right]; right++);
			Z[i] = right - left;
			right--;
		}
		else {
			int k = i - left;
			if (Z[k] < right - i + 1) {
				Z[i] = Z[k];
			}
			else {
				left = i;
				for (; right < L && str[right - left] == str[right]; right++); Z[i] = right - left; right--; } } } } string s; int main() { cin >> s;
	int L = s.length();
	string t = s;
	
	reverse(t.begin(), t.end());
	
	vector<int> Zs(L), Zt(L);
	
	Z_algorithm(s, Zs);
	Z_algorithm(t, Zt);
 
	ll ans = 0;
 
	repl(i, L / 2, L - 1) {
		// s[i]とs[i+1]の間で分割
		int a = min(L - i - 2, Zs[i + 1]);
		int b = min(L - i - 2, Zt[L - i - 1]);
		ans += max(0, a + b - (L - i - 1) + 1);
	}
 
	cout << ans << endl;
 
	return 0;
}

Zアルゴリズムの実装の詳細はググってください。このコードを読みながら紙で動作をシミュレーションしてみても結構理解できると思います。

さて、今回は以上となります。

この問題はSuffix Arrayなどを用いても解けるようですが、そちらの方はなかなか理解できなかったので、理解できたZアルゴリズムの方を採用しました。

サフィックスアレイ(接尾辞配列)については蟻本↓にも載っているので、しっかり読んで勉強して、この問題の別解も用意したいなあと思います。それではまた。

スポンサーリンク
レクタングル(大)
レクタングル(大)

シェアする

  • このエントリーをはてなブックマークに追加

フォローする