Aizu Online Judge(AOJ)が提供している「アルゴリズムとデータ構造入門」(ALDS1)の5_D問題をC++で解いてみました。
ALDS1では、「アルゴリズムとデータ構造の基礎を体得します。」とあります。トピック#5「分割統治法」では、部分的な小さな問題を解くことによって、与えられた元の問題を解く手法を紹介します。
問題(5_D: The Numbers of Inversions)
問題はリンク先をご覧ください。
AOJ ALDS1 5_D問題: The Numbers of Inversions
反転数を求めます。
反転数の計算量について
数列 $A_i$ に対して、$a_i > a_j$ かつ $i < j$ である $(i, j)$ の個数を反転数と言います。この反転数を求める問題です。
素朴なプログラムは以下となります。計算量は、数列の長さ $N$ とすると $O(N^2)$ となります。
#include <iostream>
#include <vector>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
ull result = 0;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j < n; ++j) {
if (a[i] > a[j]) {
++result;
}
}
}
cout << result << endl;
return 0;
}
このプログラムは、TLE(Time Limit Exceeded)判定となります。
$1 \leqq n \leqq 200000$ であるため、$n^2 = 4 \times 10^{10}$ となります。問題の時間制限は1秒です。1秒間に回せるループ回数は、$10^8$ 程度なので時間切れとなります。
計算量が $O(n \log n)$ となる方法を紹介します。
- セグメント木またはBIT(Bynary Index Tree)を使う(この記事で紹介するプログラム)。
- マージソートを応用する(このコースの想定解と思われる)。
セグメント木の紹介
この記事では、セグメント木の解法を紹介します。セグメント木は、以下のことができます。
- a[0]からa[n-1]の初期値を0とする。
- 任意の a[i] の値を計算量 $O(\log n)$ で増やすことができる。
- a[left] から a[right] までの合計を計算量 $O(\log n)$ で求めることができる。
具体的なセグメント木の実装とこれ以上詳細な解説については、ネットで検索するなどしていただければと思います。
このセグメント木を使うと、0, 1, 2, 3, …, n-1 を要素に持つ配列 ai の値に対して、配列の先頭から以下を計算することによって反転数 result を求めることができます。
- result = 0 として、i = 0, 1, 2, …, n-1 に対して、以下を行う。
- a[i] + 1 から n-1 までの値を result に加える。
- a[i] の値を 1 にする。
これで反転数が求めることができるのは、a[i] + 1 から n-1 までの範囲に、すでに格納されている値の和が、a[iより前] > a[i] となる個数となるためです。
C++ プログラム例(ALDS1 5_D)
セグメント木は、個人的で整備しているライブラリを流用しました(9ー47行目)。
配列 a[i] の値の範囲が $0 \leqq a_i \leqq 10^9$ と大きいため、大きさの順番を変えないように map b に格納しました(60ー65行目、座標圧縮と呼ばれる変換です)。
以下が、最終的なC++プログラムとなります。
#include <iostream>
#include <vector>
#include <set>
#include <map>
using namespace std;
typedef unsigned long long int ull;
class SegmentTree {
private:
vector<int> data;
public:
int size;
SegmentTree(int n, int init) {
size = 1;
while (size < n) {
size *= 2;
}
data.resize(2 * size - 1);
data.assign(2 * size - 1, init);
}
void update(int k, int a) {
k += size - 1;
data[k] = a;
while (k > 0) {
k = (k - 1) / 2;
data[k] = data[k * 2 + 1] + data[k * 2 + 2];
}
}
// call ST.query(a, b, 0, 0, size);
int query(int a, int b, int k, int l, int r) {
if ((r <= a)||(b <= l)) {
return 0;
}
if ((a <= l)&&(r <= b)) {
return data[k];
} else {
int vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
int vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
return vl + vr;
}
}
};
int main()
{
int n;
cin >> n;
vector<int> a(n);
set<int> s;
for (int i = 0; i < n; ++i) {
cin >> a[i];
s.insert(a[i]);
}
map<int, int> b;
int index = 0;
for (auto e : s) {
b[e] = index;
++index;
}
SegmentTree st(index, 0);
ull result = 0;
for (int i = 0; i < n; ++i) {
result += st.query(b[a[i]] + 1, index, 0, 0, st.size);
st.update(b[a[i]], 1);
}
cout << result << endl;
return 0;
}
上記プログラムは、AOJ で「AC(Accepted=正解)」と判定されました。
最後に
反転数の計算を効率よくするために、セグメント木を紹介することができました。セグメント木は、わたしにとっては難しいデータ構造ですが、使う頻度を増やして慣れていきたいです。
引き続き、ALDS1 の問題を紹介していきます。