この記事は公式解説の追加問題に関する補足解説です。以下の問題Aと問題Bの \(O(N)\) 解法を前提とします。

問題A

正整数 \(p\) と長さ \(N\) の整数列 \(A\) が与えられます。\(A\) の連続部分列のうち、バランスの良い数列であって、各整数がちょうど \(p\) 回ずつ現れるものの個数を求めてください。

問題B

正整数 \(q\) と長さ \(N\) の整数列 \(A\) が与えられます。\(A\) の連続部分列のうち、バランスの良い数列であって、ちょうど \(q\) 種類の整数が現れるものの個数を求めてください。

これを踏まえて、次の追加問題について考えます。

追加問題

長さ \(N\) の整数列 \(A\) が与えられます。\(A\) の連続部分列のうち、バランスの良い数列であるものの個数を求めてください。

\(A\) の連続部分列 \(X\) がバランスの良い数列であるとは、次の条件と同値です。

• ある \(2\) つの正整数 \(p, q\) が存在して、
  • \(X\) に含まれる整数は \(q\) 種類であり、
  • それぞれの出現回数は \(p\) である

ここで適切な整数 \(T=\Theta(\sqrt N)\) を選び、次の \(2\) つの場合に分けて、個別に数え上げを行います。

• \(p\leq T\) の場合
• \(p\gt T\) の場合

---

\(p\leq T\) の場合、\(p\) を固定すると問題は「各整数が \(p\) 回ずつ現れるバランスの良い数列の個数」を数えることになり、問題Aと同様に解くことができます。

したがって、問題Aの \(O(N)\) 解法を各 \(p\in\{1, \dots, T\}\) に適用すると \(O(NT)=O(N\sqrt N)\) 時間で計算できます。

---

\(p\gt T\) の場合を考えます。バランスの良い数列の長さは \(pq\leq N\) を満たす必要があるため、

\[q\leq\frac Np\lt\frac NT=O(\sqrt N)\]

となります。したがって \(q\) の値はたかだか \(O(\sqrt N)\) 通りであり、それぞれについて処理すればよいです。

\(q\) を固定すると、問題は「\(q\) 種類の整数が現れるバランスの良い数列の個数」を数えることになります。これは本質的に問題Bと同じですが、\(p\gt T\) という制約を追加で考慮する必要があります。

連続部分列 \((A_\ell, \dots, A_{r-1})\) の右端 \(r\) を固定します。\(p\gt T\) という条件を一旦無視すれば、ちょうど \(q\) 種類の整数が現れるような左端 \(\ell\) は区間になっており、区間の下限と上限をそれぞれ \(d_r, u_r\) として区間を \([d_r, u_r)\) で表すことができます。特に、右端の整数 \(A_{r-1}\) から左向きに見ていったとき、異なる整数の種類数がちょうど \(q\) になる最初の位置は \(u_r-1\) であり、さらに（存在するなら）\(q+1\) 種類になる最初の位置は \(d_r-1\) です。

ここで \(p\gt T\) の追加条件を考えると、バランスの良い数列の長さに関して

\[r-\ell=pq\gt Tq\]

が成り立つため、\(\ell\lt r-Tq\) が必要になります。したがって左端の上限を \(u_r^\prime\leftarrow\min\{u_r, r-Tq\}\) で置き換えて、各 \(r\) に対し \(\ell\in[d_r, u_r^\prime)\) を探索すれば、「\(q\) 種類の整数が \(T+1\) 回以上現れるバランスの良い数列」を数えることができます。

公式解説と同様に、右端 \(r\) を昇順に更新して、整数の集合が同じなら累積和を連想配列で保持して、変わるなら再計算する方針を取ります。集合が変わるとき、\(A_{u_r-1}\) は新しい集合に含まれないため、右端 \(r+1\) に対する左端の下限を \(d_{r+1}\) として、\(u_r\leq d_{r+1}\) が成り立ちます。また、\(u_r^\prime=\min\{u_r, r-Tq\}\leq u_r\) であるため、\(u_r^\prime\leq u_r\leq d_{r+1}\) となります。このことから、各左端 \(\ell\) について、実際に累積和の処理対象となる集合はたかだか \(1\) つに限られます。よって、各 \(\ell\) に対する累積和の計算は全体で最大 \(N\) 回になり、\(1\leq q\lt \frac NT\) を全探索すると \(O(N\sqrt N)\) 計算量となります。

---

以上より、追加問題を \(O(N\sqrt N)\) 時間で解くことができます。