問題
公式の問題はこちらです．
問題文
$N$ 個の整数 $A_1, …, A_N$ が与えられます．$1 ≦ i < j ≦ N$ を満たす全ての $i, j$ の組についての $|A_i – A_j|$ の総和を求めてください．
制約
- $2 ≦ N ≦ 2 \times 10^5$
- $|A_i| ≦ 10^8$
- $A_i$ は整数である．
解説
はじめに，TLE（Time Limit Exceeded）してしまった「ボツ案」を紹介し，その後計算量を減らしたアルゴリズムを解説します．
ボツ案（TLEした解法）
この問題，与えられた整数から任意の二つの整数 $A_i, A_j$ $(1 ≦ i < j ≦ N)$ を取り出して，その二つの整数の差分の絶対値 $|A_i – A_j|$ を足していけばよいのでは？と最初考えました．
$A_i, A_j$ $(1 ≦ i < j ≦ N)$ が取りうるすべての組み合わせを抽出するため，このサイトで紹介されているアルゴリズムを使いました．詳細は把握できていませんが，recursive_comb という関数を再帰的に使って，全通りの組み合わせを出力できます．
この考え方で書いたのが以下のソースコードです．
#include <bits/stdc++.h>
using namespace std;
#include <functional>
#include <cstdlib> // abs() for integer
void recursive_comb(int *indexes, int s, int rest, std::function<void(int *)> f) {
if (rest == 0) {
f(indexes);
} else {
if (s < 0) return;
recursive_comb(indexes, s - 1, rest, f);
indexes[rest - 1] = s;
recursive_comb(indexes, s - 1, rest - 1, f);
}
}
void foreach_comb(int n, int k, std::function<void(int *)> f) {
int indexes[k];
recursive_comb(indexes, n - 1, k, f);
}
int ans = 0;
int A[200010];
int main(void) {
int N;
cin >> N;
for (int i=0; i<N; i++) cin >> A[i];
// calculate absolute values for all the combination.
foreach_comb(N, 2, [](int *indexes) {
ans += abs(A[indexes[0]] - A[indexes[1]]);
});
cout << ans << endl;
return 0;
}全通りの組み合わせを考えると，$_NC_2 = \frac{N(N-1)}{2}$ 回の計算が最低限必要になる（アルゴリズムの実装によってはもっと計算量が増える）ので，これじゃTLEしちゃうかなと思っていましたが，案の定TLEしました (~.~;
ACした解法（ソート + 累積和）
この問題では任意の $i, j$ に対して $|A_i – A_j|$ を計算しますが，絶対値を外して計算を簡略化することを考えます．ソートをして絶対値を外しましょう．
昇順でソートをすると，任意の $1 ≦ i < j ≦ N$ に対して，$|A_i – A_j| = A_j – A_i$ が成り立つので，条件を満たす任意の $i, j$ に対して $A_j – A_i$ の総和を考えればよく，絶対値を外すことができます．また，ソートをして絶対値を取っても，答えには影響ありません．
このとき，答えは以下のようになります．
$\displaystyle \sum_{j=2}^{N} (A_j – A_{j-1}) + (A_j – A_{j-2}) + … + (A_j – A_2) + (A_j – A_1) $
$ = \displaystyle \sum_{j=2}^{N} (j – 1) A_j – (A_{j-1} + A_{j-2} + … + A_2 + A_1)$
この $A_{j-1} + A_{j-2} + … + A_2 + A_1, (2 ≦ j ≦N)$ の部分が累積和 $S_j$ となります．この部分はあらかじめ計算しておいて，配列 S[j] に格納しておくことで，計算量を抑えることができます．
累積和を配列に格納する部分，上記のシグマの計算部分はいずれも $O(N)$ なので，計算量を $O(N)$ に抑えられたことになります．なお，ソートの部分がボトルネックとなり，全体の計算量は $O(N \log N)$ となりますが，元々の $O(N^2)$ よりは計算量が抑えられたことになります．
累積和に関してはこのサイト，この問題の解き方に関してはけんちょんさんの記事を見て勉強させてもらいました．このような累積和の概念は，$N ≒ 10^5$ の制約の問題ではよく使うみたいです．
補足考察
計算量についてはこの記事でも考察しましたが，やっぱりちゃんと考えないとダメですね．今回は試しに，$O(N^2)$ と $O(N \log N)$ でどれくらい計算量が違うのかを考えてみます．（今回の問題では，ソートがボトルネックで計算量が $O(N \log N)$ となっています．）
今回の例だと，$N ≦ 2 \times 10^5$ なので，この条件下の最大値で以下の計算をしてみました．
- $N^2 = 40000000000 = 4.0 \times 10^{10}$
- $N \log N ≒ 1060206 ≒ 1.1 \times 10^6$
※上記の計算で $\log$ の底は $10$ としていますが，底がなんであろうと定数倍の差にしかならないので計算量を考える場合は特段気にする必要はなさそうです．（参考）
$N$ が $10^5$ のオーダーだと，$O(N^2)$ と $O(N \log N)$ に大きな差があることがわかりますね．
ソースコード
下記のソースコードでACしました．
#include <bits/stdc++.h>
using namespace std;
int main(void) {
int N;
cin >> N;
vector<long long> A(N+1, 0);
vector<long long> S(N+1, 0);
for (int i=1; i<=N; i++) cin >> A[i];
// |Aj - Ai| equals to (Aj - Ai) by sorting A array. O(N log N)
sort(A.begin()+1, A.end());
// calculate cumulative sum, O(N)
for (int i=2; i<=N; i++) { S[i] = A[i-1] + S[i-1]; }
// calculate answer, O(N)
long long ans = 0;
for (int j=2; j<=N; j++) { ans += (j-1)*A[j] - S[j]; }
cout << ans << endl;
return 0;
}問題文と添字をそろえるため，A[1] から値を格納していきました．ソートの部分で sort(A.begin(), A.end()); としたら，A[0] = 0 もソートの範囲に含まれてしまい，WAになってしまいました．そこで，値が格納されている範囲のみでソート（sort(A.begin()+1, A.end());）をして，ACとすることができました！
与えられた $A_i, (i \in \{1, …, N\})$ に負の数が含まれていたときに，不具合が起きたのだと思われます．
まとめ・感想
今回の問題に関しては，累積和を使うパターンとして記憶しておく必要がありそうですね．$N ≒ 10^5$ の制約の問題を見たら反射的に $O(N^2)$ ではTLEだと考えるべきなのかもしれません．
コメント