公式解説と同様に \(A_i\) より小さい要素が \(A_i\) のすぐ左にいくつ隣接するかを \(L_i\) 、右についても同様に \(R_i\) とします。

このとき \(A\) の連続部分列のうち長さが \(k\) のものについて答えが \(A_i\) になるようなものの個数は

\(\displaystyle [x^k]x(1+x+x^2+\dots+x^{L_i})(1+x+x^2+\dots+x^{R_i})=[x^{k-1}]\frac{(1-x^{L_i+1}) (1-x^{R_i+1})}{(1-x)^2}\) 個あります。

したがって、長さ \(N+2\) の配列 \(s=(s_0,s_1,\dots,s_{N+1})\) を用意し、\(i=1,2,\dots,N\) について

• \(s_0\gets s_0+A_i\)
• \(s_{L_i+1}\gets s_{L_i+1}-A_i\)
• \(s_{R_i+1}\gets s_{R_i+1}-A_i\)
• \(s_{L_i+R_i+2}\gets s_{L_i+R_i+2}+A_i\)

とした後、\(s\) の累積和を取る処理を \(2\) 回行うと

\(A\) の連続部分列のうち長さが \(k\) のものについての答えは \(s_{k-1}\) となります。

計算量は \(L_i,R_i\) を求めるのは公式解説と同様に \(O(N\log N)\)、配列 \(s\) を計算するのは \(O(N)\) より、\(O(N\log N)\) となります。

実装例 ( C++ , 321 ms )