https://atcoder.jp/contests/abc339/editorial/9204

kyopro_friends さんによるこちらの解説で平方分割解法が紹介されています。

似たような考えで計算量が改善できて、 \(O((Q+N)\log(N) + (QN)^{\frac{2}{3}} )\) になります。

計算量は悪化しているように見えますが、実際にはほとんどの提出より実行時間が短いです。

以降、 \(A\) などの整数列は 0-index とします。

[0] 定義

まず、以下の条件を満たす \((0,1,\dots ,N - 1)\) 順列 \(P=(P_{0},P_{1},\dots ,P_{N-1})\) を \(1\) つとります。

• 任意の整数 \(i,j\) について、\(P_{i}<P_{j}\) ならば \(A_{i}<A_{j}\)

そして、関数 \(f(r,a)\) を以下のように定義します。

• \(i=0,1,\dots ,r-1\) のうち、 \(P_{i}<a\) を満たす \(i\) についての \(A_{i}\) の総和

また、 \(B\) をバケットサイズ、つまり正整数とします。

[1] 概要

\(A\) の中で \(X\) 以下であるような値の数を \(Y\) とすると、答えは \(f(R,Y) - f(L-1,Y)\) と表せます。\(Y\) は \(A\) をソートしたものを用いて \(O(\log(N))\) で求まります。よって、\(f(r,a)\) が高速に求まれば答えも高速にも求まります。

クエリごとに \(f(r,a)\) を高速に求めるために、全ての \(r=0,1,\dots, \left\lfloor\frac{N-1}{B}\right\rfloor\) と \(a=0,1,\dots ,\left\lfloor\frac{N-1}{B}\right\rfloor\) について、 \(f(Br,Ba)\) を前計算します。

そして、 \(f(r,a) - f(B\left\lfloor\frac{r}{B}\right\rfloor,B\left\lfloor\frac{a}{B}\right\rfloor)\) をクエリごとに求めます。

[2] 前計算

\(f(r,a)\) とは \(i<r\) かつ \(P_{i}<a\) を満たす \(i\) に対する \(A_{i}\) の総和です。よって累積和をとると、空間・時間計算量ともに \(O(N^{2})\) で、クエリごとに \(O(1)\) で求めることができます。この累積和の取り方は区切りが \(1\) ごとでしたが、 \(B\) ごとにすることで、空間・時間計算量ともに \(O(\frac{N^2}{B^2})\) で、\(f(B\left\lfloor\frac{r}{B}\right\rfloor,B\left\lfloor\frac{a}{B}\right\rfloor)\) を \(O(1)\) で求めることができます。

具体的には以下のようにします。

要素が全て \(0\) で初期化された \(\left(\left\lfloor\frac{N-1}{B}\right\rfloor+1\right)\times \left(\left\lfloor\frac{N-1}{B}\right\rfloor+1\right)\) の二次元配列 \(S\) に対して、\(i=0,1\dots, N-1\) の順に以下の操作を行います。

• \(S[\left\lfloor\frac{i}{B}\right\rfloor+1][\left\lfloor\frac{P_{i}}{B}\right\rfloor+1]+=A_{i}\)

その後累積和をとる操作をします。

これを行うと \(S[r][a]=f(Br,Ba)\) が成り立ちます。

[3] クエリごとの計算

\(f(r,a) - f(B\left\lfloor\frac{r}{B}\right\rfloor,B\left\lfloor\frac{a}{B}\right\rfloor)\) は以下の \(2\) つの条件のいずれかを満たす \(i\) についての \(A_{i}\) の総和です。

• \(B\left\lfloor\frac{r}{B}\right\rfloor \leq i< r \) かつ \(P_{i}<a\)
• \(B\left\lfloor\frac{a}{B}\right\rfloor \leq P_{i}<a\) かつ \(i<B\left\lfloor\frac{r}{B}\right\rfloor\)

この \(2\) つの条件をどちらも満たす \(i\) は、\(i\) に対する不等号の制約によって存在しないため、どちらについても独立に求めてその和を取ればいいです。

\(B\left\lfloor\frac{r}{B}\right\rfloor \leq i< r \) を満たす \(i\) 全てに対して、 \(P_{i}<a\) を満たすか否か探索すれば \(1\) つめの条件を満たす \(i\) を過不足なく見つけることができます。

\(2\) つめの条件についても、\(B\left\lfloor\frac{a}{B}\right\rfloor \leq P_{i}<a\) を満たす全ての \(i\) に対して、\(i<B\left\lfloor\frac{r}{B}\right\rfloor\) を満たすか否か探索すれば良いです。

\(1\) つめの条件は愚直に、 \(2\) つめの条件は \(\text{inv}_{P_{i}}=i\) を満たす順列 \(\text{inv}\) を事前に求めておくことで、どちらも \(O(B)\) で求められます。

[4] 計算量解析

• [0] で \(P\) を求めるのに \(A\) をソートする必要があります。

• [1] でクエリごとに \(Y\) を求める必要があります。これは stl::upper_bound 等を用いて \(O(\log(N))\) で求められます。

• [2] での前計算の計算量は \(O(\frac{N^{2}}{B^{2}})\) でした。

• [3] での計算量は \(O(B)\) でした。

以上より計算量は \(O(N\log(N) + \frac{N^{2}}{B^{2}} +Q(B+\log(N)))\) です。ここで、\(\frac{N^{2}}{B^{2}} =QB\) を変形すると \(B=\left(\frac{N^2}{Q}\right)^{\frac{1}{3}}\) となります。 \(B\) をこの値に置き換えることで、初めに記述した計算量 \(O((Q+N)\log(N) + (QN)^{\frac{2}{3}} )\) を得ることができます。

[5] 実装

\(B\) は \(50\) 以上 \(100\) 以下の値にしておくのが良いでしょう。

公式解説による解法の計算量は \(O(Q\log(N)^{2}+N\log(N))\) であると推測され、これはこの解法の計算量より良い計算量でありますが、今回の制約ではこちらの解法の方も十分高速です。

実装例 c++ 231 ms

実装例 pypy 1095 ms