±1 Operation 2
問題
出典
ABC255 - D 788 diff
考察
配列A はソートしても一般性を失わないことはある程度競技プログラミングをやって来ればすぐにわかる
したがって、以下 A]は[*! ソートして昇順で扱う
次に、各クエリでのX_iがソートしたAのどこに来るのかを考える
つまり、A_1 \leq A_2 \leq ... \leq A_K \leq X_i \leq A_{K + 1} \leq ... \leq A_N を満たすKを探索する
これは Python であれば
bisect モジュールを使用すればO(\log N)すぐに求められるこの場合に必要な操作の回数を考える
すると、[*% A_1 〜 A_Kまでは X_i以下]であり、[*% A_{K + 1} 〜 A_NはX_i以上]であることがわかるので以下の式で操作回数は求められる
\{ (X_i - A_1) + (X_i - A_2) + ... + (X_i - A_K) \} + \{ (A_{K + 1} - X_i) + (A_{K + 2} - X_i) + ... + (A_N - X_i) \}
= \{ K \cdot X_i - (A_1 + A_2 + ... + A_K) \} + \{ (A_{K + 1} + A_{K + 2} + ... + A_N) - (N - K) \cdot X_i \} \\
= K \cdot X_i - \sum_{j = 1}^{K} A_j + \sum_{j = K + 1}^{N} A_j - (N - K) \cdot X_i
= \sum_{j = K + 1}^{N} A_j - \sum_{j = 1}^{K} A_j + (2K - N) \cdot X_i
= \left( \sum_{j = 1}^{N} A_j - \sum_{j = 1}^{K} A_j \right) - \sum_{j = 1}^{K} A_j + (2K - N) \cdot X_i
= \left( \sum_{j = 1}^{N} A_j - 2 \sum_{j = 1}^{K} A_j \right) + (2K - N) \cdot X_i
この\sum_{j = 1}^{K} A_jはAに対する累積和 Sであらかじめ前処理で求めておくと、各クエリに関してO(\log N)で求めることができる
したがって、各クエリは O(\log N)で求めることが出来たので、[* 計算量は O \left((N + Q) \log N \right)]となりました
コード
Pythonfrom bisect import bisect_right
n, q = map(int, input().split())
a = sorted(list(map(int, input().split())))
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = a[i] + s[i]
sum_ = sum(a)
for _ in range(q):
X = int(input())
K = bisect_right(a, X) # O(log N)
print(sum_ - 2 * s[K] + (2 * K - n) * X)