ABC255D - ±1 Operation 2

考え方

クエリごとに$X$と$A_{i}$の差を求めるとTLEする.

クエリの計算量$O(Q)$は絶対に必要なので,操作回数を高速に求める必要がある.

ここで,

  • $A_{i} < X$なら操作回数は$X - A_{i}$
  • $A_{i} \geq X$なら操作回数は$A_{i} - X$
だから,答えは
\begin{aligned}
\sum_{i (A_{i} < X)} (X - A_{i})
+ \sum_{i (A_{i} \geq X)} (A_{i} - X)
\end{aligned}
となる.

よって,$A$をソートして累積和を求めておけばよい.

回答例

from bisect import bisect_left

N, Q = map(int, input().split())
A = list(map(int, input().split()))
A.sort()
acc = [0] * (N + 1)
for i in range(N):
    acc[i + 1] = acc[i] + A[i]

for _ in range(Q):
    x = int(input())
    ind = bisect_left(A, x)
    if ind < N:
        l = acc[ind]
        r = acc[-1] - l
        print(x * ind - l + r - x * (N - ind))
    else:
        print(x * N - acc[-1])