Official
F - Double Sum Editorial by Nyaan
この問題は平面走査と呼ばれるアルゴリズムの基本問題です。この問題を解けなかった人は平面走査について学習するとよいでしょう。
求めたい二重和は
\[\sum_{i=1}^N \sum_{i < j} \max(A_j - A_i, 0)\]
です。ここで、\(i\) を固定したときの二重和への寄与は
\[ \begin{aligned} &\sum_{i < j}\max(A_j - A_i, 0) \\ &= \sum_{i < j, A_i \leq A_j} A_j - A_i \\ &= (i < j かつ A_i \leq A_j である A_j の総和) \\ &- (i < j かつ A_i \leq A_j である A_j の個数) \times A_i \end{aligned} \]
と言い換えることが出来ます。 この事実を利用すると、次の手順で平面走査をすることでこの問題を解くことが出来ます。
- 次の 2 種類の値を管理するデータ構造を用意する。
- 「要素の追加」「\(x\) 以上の要素の個数」の 2 種類のクエリを処理できる多重集合 \(S_0\)
- 「要素の追加」「\(x\) 以上の要素の総和」の 2 種類のクエリを処理できる多重集合 \(S_1\)
- また、答えを格納する変数 \(\mathrm{ans}\) を用意する。はじめ \(\mathrm{ans} = 0\) とする。
- \(i = N, N-1, \dots, 2, 1\) の順に次の処理を行う。
- \(c\) を \(S_0\) に \(x = A_i\) でクエリを投げたときの返り値とする。
- \(s\) を \(S_1\) に \(x = A_i\) でクエリを投げたときの返り値とする。
- \(\mathrm{ans}\) に \(s - c \times A_i\) を加算する。
- \(S_0, S_1\) に \(A_i\) を追加する。
- 最終的な \(\mathrm{ans}\) の値を出力する。
\(S_0, S_1\) については、座標圧縮した Fenwick Tree を持てば 1 回の処理あたり \(\mathrm{O}(\log N)\) の計算量で処理出来ます。
以上よりこの問題を \(\mathrm{O}(N \log N)\) で解くことが出来て、これは十分高速です。
- 実装例 (Python)
from atcoder.fenwicktree import FenwickTree
import bisect
N = int(input())
A = list(map(int, input().split()))
B = sorted([x for x in set(A)])
M = len(B)
sum0 = FenwickTree(M)
sum1 = FenwickTree(M)
ans = 0
for i in reversed(range(N)):
k = bisect.bisect_left(B, A[i])
c = sum0.sum(k, M)
s = sum1.sum(k, M)
ans += s - c * A[i]
sum0.add(k, 1)
sum1.add(k, A[i])
print(ans)
posted:
last update: