G - Balanced Subarrays Editorial
by
vwxyz
原案:vwxyz
まず、\(1\) つ目の値(各数字が \(B_k\) 回ずつ登場するもの)について考えます。
連続部分列の右端をすべて試すことにします。
このとき、同じ要素が \(B_k+1\) 個以上現れないような左端の下限が存在します。
これは尺取り法などにより、各右端について合計 \(O(N)\) で求めることができます。
左端がこれより左にないとき、条件は、すべての要素の登場回数が \(\bmod B_k\) で \(0\) に等しいことと言い換えられます。
十分大きな素数 \(M\) を取ります。
各値ごとに総和が \(M\) の倍数になるような \(B_k\) 個の乱数を用意し、先頭から見て各要素の登場回数\(\pmod {B_k}\) に応じた乱数を割り当てて、各連続部分列のハッシュ値をその総和\(\pmod {B_k}\) とすると、すべての要素の登場回数が \(\bmod B_k\) で \(0\) に等しいことはハッシュ値が \(0\) になることと十分高い確率で同値です。
よって、ハッシュ値の累積和をとっておき、ハッシュ値をキーとした連想配列でインデックスを持っておくことにより、\(O(N)\) や \(O(N \log N)\) で総数を求めることができます。
次に、\(2\) つ目の値(\(B_k\) 種類の数字が登場するものもの)について考えます。
連続部分列の右端をすべて試すことにします。
ちょうど \(B_k\) 種類の要素が現れるような左端は区間になっています。
十分大きな素数 \(M\) を取ります。 同じ値には同じになるように、各要素にハッシュとして乱数を割り振り、連続部分列(または集合)のハッシュ値をそれに含まれるハッシュの総和を \(M\) で割った余りとします。
先頭から \(i(0 \leq i \leq N)\) 個のハッシュ値の総和を \(S_i\) と表すことにします。
また、右端を固定したとき、ちょうど \(B_k\) 種類の要素が現れるような部分列に含まれる要素は一意に定まるため、これを \(1\) つずつ含むような集合のハッシュ値を \(H\) とします。
すると、\(A\) の \(l+1\) 番目から \(r\) 番目までの連続部分列にすべての要素が同じ回数ずつ現れることは
\(S_r-S_l \equiv H \times \frac{r-l}{B_k} \pmod{M}\)
と十分高い確率で同値になります。
これを変形すると、
\(S_r \times B_k - H \times r=S_l \times B_k - H \times l\)
となり、区間内の各 \(i\) に対して \(S_i \times B_k - H \times i\) の頻度が管理できれば条件を満たす左端の個数を求めることができます。
右端を左から右に動かしていったとき、\(B_k\) 種類の要素の集合が同じならばこれを継続して管理することにし、集合が変わるならば新しく計算し直すことにします。
集合が変わるとき、左端の下限は前回の左端の上限よりも大きいため、各 \(i\) に対して高々 \(1\) 回しか \(S_i \times B_k - H \times i\) を計算しないことがわかります。
適切に実装することで、\(O(N)\) や \(O(N \log N)\) で総数を求めることができます。
追加問題
バランスの良い数列の個数を \(O(N \sqrt N)\) で求める方法を考えてみましょう。
posted:
last update:
