D - ±1 Operation 2 Editorial by en_translator
First, sort \(A\) so that \(A_1 \le A_2 \le \dots \le A_N\).
When \(A_1 \le A_2 \le \dots \le A_{k} \le X_i \le A_{k+1} \le \dots \le A_N\), the answer is expressed as
\((X_i-A_1) + (X_i-A_2) + \dots + (X_i-A_k) + (A_{k+1}-X_i) + \dots +(A_N-X_i) = \sum^{k}_{j=1} (X_i-A_j) + \sum^{N}_{j=k+1} (A_j-X_i).\)
\(\sum^{k}_{j=1} (X_i-A_j) = kX_i - \sum^{k}_{j=1} A_j\), whose \(kX_i\) can obviously be calculated in an \(O(1)\) time, and \(\sum^{k}_{j=1} A_j\) can be calculated in an \(O(1)\) time for each query by precalculating the cumulative sum in an \(O(N)\) time. Also, \(\sum^{N}_{j=k+1} (A_j-X_i)\) can be calculated in the similar way.
Additionally, we can find \(k\) such that \(A_k \le X_i \le A_{k+1}\) with binary search in \(O(\log N)\) time per query.
Therefore, the problem has been solved in a total of \(O(N + Q \log N)\) time.
Sample code (C++):
#include<bits/stdc++.h>
using namespace std;
int main(){
int n,q;
cin >> n >> q;
vector<long long> a(n);
for(auto &nx : a){cin >> nx;}
sort(a.begin(),a.end());
vector<long long> rw(n+1,0);
for(long long i=0;i<n;i++){rw[i+1]=rw[i]+a[i];}
for(int i=0;i<q;i++){
long long x;
cin >> x;
int st=0,fi=n-1;
while(st<=fi){
int te=(st+fi)/2;
if(a[te]<x){st=te+1;}
else{fi=te-1;}
}
long long res=x*st;
res-=rw[fi+1];
res+=(rw[n]-rw[st]);
res-=x*(n-st);
cout << res << '\n';
}
return 0;
}
posted:
last update: