Official

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: