Official

A - 学校対抗マラソン大会 / Inter-School Marathon Competition Editorial by physics0523


初心者の方へ

愚直に二重ループを回す解法もありますが、高速な解法を示します。

ある数列 \(A_i\) に対して \(\sum_{i<j} |A_i-A_j|\) を求められればこの問題に正解できます。

数列の長さを \(N\) として、 \(O(N^2)\) かけて二重ループを回してもよいですが、以下の方法で高速に求めることができます。

  • まず、 \(A\) を昇順にソートする。ソート後の列について、 \(A=(A_1,A_2,\dots,A_N)\)
  • これまでの総和 \(s=0\) を保持する。
  • その後、 \(i=1,2,\dots,N\) の順に以下を行う。
    • 答えに \((i-1) \times A_i - s\) を加算する。
    • その後、 \(s\) に \(A_i\) を加算する。

\(A\) を昇順ソートすることで、求める値を絶対値記号を外して \(\sum_{i<j} A_j-A_i\) とすることができます。
ループの \(A_k\) を参照する際に、 \(A_k\) が左側に来る項を全て足し上げることを考えます。
\((A_k-A_1) + (A_k-A_2) + \dots + (A_k-A_{k-1})\) を求めたいのですが、これを式変形すると \(A_k \times (k-1) - (A_1+A_2+\dots+A_{k-1})\) となり、ループで \(A_k\) を見るタイミングで確かにこの値が加算できていることが分かります。

あとは、同じ学校同士でこれをまとめて計算すればよいです。これはバケットソートの要領で実現可能です。
時間計算量は \(O(N \log N)\) です。

Bonus! \(s\) を保持せずに計算することもできます。どのようにすればよいでしょうか?

#include<bits/stdc++.h>

using namespace std;
using ll=long long;

ll f(vector<ll> &v){
  sort(v.begin(),v.end());
  ll res=0,pre=0,pc=0;
  for(auto &nx : v){
    res+=(nx*pc-pre);
    pre+=nx; pc++;
  }
  return res;
}

int main(){
  ll N;
  cin >> N;
  vector<vector<ll>> v(N+1);
  for(ll i=0;i<N;i++){
    ll C,L;
    cin >> C >> L;
    v[C].push_back(L);
  }
  ll res=0;
  for(auto &nx : v){
    res+=f(nx);
  }
  cout << res << "\n";
  return 0;
}

posted:
last update: