公式解説と同様にして、以下の問題を解きます。

\(i=N, N+1, \cdots, 2N-1\) について、以下の答えを足し合わせよ。

• \(R_{i-N+1}, R_{i-N+2}, \cdots, R_{i-1}\) の中に、 \(R_i\) と等しい要素は何個ありますか？

準備

まず、各 \(x=0,1,\cdots, M-1\) について、 \(x\) に関する空のリスト \(\mathrm{bucket}(x)\) を用意します。

次に、 \(i=0,1,\cdots, 2N-1\) の順に、 \(i\) を \(\mathrm{bucket}(R_i)\) の末尾に追加します。

これで、各 \(\mathrm{bucket}(x)\) は、 \(R_i = x\) となる \(i\) が小さい順に入っています。

問題に答える

次のクエリに高速に答えます。

• \(R_{i-N+1}, R_{i-N+2}, \cdots, R_{i-1}\) の中に、 \(R_i\) と等しい要素は何個ありますか？

これは次のように言い換えることができます。

• \(\mathrm{bucket}(R_i)\) の中で、 \((i-N+1)\) 以上 \((i-1)\) 以下の要素は何個ありますか？

これは以下のように答えることができます。

• \(r\) は \(\mathrm{bucket}(R_i)\) の中で \((i-1)\) 以下の要素の個数とする。
• \(l\) は \(\mathrm{bucket}(R_i)\) の中で \((i-N)\) 以下の要素の個数とする。
• \(\mathrm{bucket}(R_i)\) の中で \((i-N+1)\) 以上 \((i-1)\) 以下の要素は \((r-l)\) 個である。

\(\mathrm{bucket}(x)\) は小さい順に値が入っているので、二分探索を用いて \(r, l\) を高速に求めることができ、各クエリに \(O(\log N)\) 時間で答えることができます。

したがって、この問題は全体で時間計算量 \(O(M + N \log N)\) で解くことができます。

公式解説は時間計算量が \(O(M+N)\) なので、公式解説の解法の方がより効率的ですが、いま紹介した解法でも十分高速に AC することができます。

shobonvipの実装 (C++, 98ms)

shobonvipの実装 (Pypy3, 282ms)