D - ±1 Operation 2 Editorial by physics0523
まず、 \(A\) をソートして \(A_1 \le A_2 \le \dots \le A_N\) としておきましょう。
\(A_1 \le A_2 \le \dots \le A_{k} \le X_i \le A_{k+1} \le \dots \le A_N\) であるとき、答えは以下の式で表されます。
\((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\) となり、このうち \(kX_i\) は明らかに \(O(1)\) 、 \(\sum^{k}_{j=1} A_j\) は累積和を用いて事前計算 \(O(N)\) 、クエリあたり \(O(1)\) で求めることが出来ます。また、 \(\sum^{N}_{j=k+1} (A_j-X_i)\) についても同様に求めることが出来ます。
また、 \(A_k \le X_i \le A_{k+1}\) となるような \(k\) は二分探索によりクエリあたり \(O(\log N)\) で求めることが出来ます。
よって、この問題を \(O((N + Q) \log N)\) で解くことが出来ました。
実装例(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: