ARC099 F : Eating Symbols Hard(1200)
作成日:2018.10.03
最終更新日:2018.10.03
問題概要
$-10^9$から$10^9$までの添え字が付いた長さ$2 \times 10^9 + 1$の整数列$\{A_i\}$と整数$P$がある。
初め、この数列の各要素と$P$は$0$である。
記号$+,-,\gt,\lt$はそれぞれ「$A_P$を$1$大きくする」、「$A_P$を$1$小さくする」、「$P$を$1$大きくする」、「$P$を$1$小さくする」という意味を持つ。
記号列$S$の$i$文字目から$j$文字目を処理した後の整数列が、$1$文字目から$N$文字目を処理した後の整数列と等しくなるような組$(i,j)$$(1 \leq i \leq j \leq N)$の個数を求めよ。
制約
$1 \leq N \leq 2.5 \times 10^5$
$|S| = N, \; S_i = {'+'},{'-'},{'\gt'},{'\lt'}$
考え方
$S$の$i$文字目から$j$文字目を処理した後の整数列を$A[i,j]$とする。
但し、$i \gt j$の時は$A[i,j]$の各要素は$0$であるとする。
組$(i,j)$ごとに$A[i,j]$を求めるわけにはいかないから、累積和のようにしてこれを求められないかを考える。
単に$A[1,j]-A[1,i-1]$としても、$S$を初めから$i$文字目まで処理した後の$P$の値($P_i$とする)が$0$でなければ等しくない。
しかし逆に考えると、$A[i,j]$と$A[1,j]-A[1,i-1]$の違いは各要素の位置が$P_{i-1}$だけずれているだけである。
そこで位置のずれた要素が等しいかを高速に判定するために、Rolling Hashを用いることにする。
具体的には、数列$\{A_i\}$を$\sum b^i A_i$($= H(A)$とする)と変換する。
すると要素の位置を$x$だけ後ろにずらすのは$b^x$倍することで実現できるから、$H(A[i,j])$$= (H(A[1,j])$$- H(A[1,i-1]))$$b^{-P_{i-1}}$である。
Hashが衝突しないとすると、求めるべきは$H(A[1,N])$$= H(A[i,j])$となる組$(i,j)$の数である。
これを愚直に試すことはできないので工夫する。
$i$を固定すると、条件は$H(A[1,j])$$= H(A[1,N]) b^{P_{i-1}}$$+ H(A[1,i-1])$となる。
このような$j$の個数は$H(A[1,j])$をmapで管理すれば$O(log \: N)$で求められる。
すべての$i$について、これを求めて足し合わせたものが答えとなる。
$H$をそのまま持つわけにはいかないから、素数で割った余りをいくつか持っておくことにする。
また、$b$は$A_0$と$b A_1$の区別がつくように$b \gt N$であるべきである。
Rolling Hashの構築の計算量は$O(N)$であり、全体の計算量は$O(N \; log \: N)$である。
これは十分高速である。
ソースコード
マクロ等はこちら
LL inv(LL a , LL p){
LL ret = 1;
for(LL i = p-2 ; i ; i /= 2){
if(i&1) (ret *= a) %= p;
a = a * a % p;
}
return ret;
}
const int RH_prime = 3;
const LL RH_mod[] = {1000000007,1000000021,1000000033};
const LL base = 314159;
struct RH{
int n;
vector<LL> hash[RH_prime];
vector<LL> pw[RH_prime];
RH(){n=0;}
RH(const string &s){
n = (int)s.size();
for(int i = 0 ; i < RH_prime ; ++i){
hash[i].push_back(0);
pw[i].push_back(1);
for(int j = 0 ; j < n ; ++j){
if(s[j] == '+' || s[j] == '-'){
hash[i].push_back((hash[i][j] + (s[j]=='-'?RH_mod[i]-pw[i][j]:pw[i][j])) % RH_mod[i]);
pw[i].push_back(pw[i][j]);
} else {
hash[i].push_back(hash[i][j]);
pw[i].push_back(pw[i][j] * (s[j]=='>'?base:inv(base,RH_mod[i])) % RH_mod[i]);
}
}
}
}
};
int main(){
int N; cin >> N;
string S; cin >> S;
RH rh(S);
map<tuple<LL,LL,LL>,int> M;
LL ans = 0;
repm(i,N,0){
++M[make_tuple(rh.hash[0][i],rh.hash[1][i],rh.hash[2][i])];
tuple<LL,LL,LL> z = make_tuple((rh.hash[0][N]*rh.pw[0][i-1]+rh.hash[0][i-1])%RH_mod[0],
(rh.hash[1][N]*rh.pw[1][i-1]+rh.hash[1][i-1])%RH_mod[1],
(rh.hash[2][N]*rh.pw[2][i-1]+rh.hash[2][i-1])%RH_mod[2]);
ans += M[z];
}
cout << ans << endl;
return 0;
}
解法まとめ
数列$\{A_i\}$を$\sum b^i A_i$と変換するRolling Hashを考える。
構築は文字列を先頭から見ていき、$+$を処理したら$b^P$、$-$を処理したら$-b^P$だけ足せば良い。
その際、$P$の代わりに$b^P$も計算していく。(21-34行)
$i$番目のRolling Hashの値を$H_i$とする。
$i$を$N$から$1$まで降順に見ていき、mapに$H_i$を詰めた後にmap内にある$H_N b^{P_{i-1}}$$+ H_{i-1}$の数を見る。
この数の総和が答えである。(43-51行)