1. はじめに
この記事では、転倒数(反転数)と呼ばれる問題を BIT を使って解く。転倒数の詳細は次の通り。
長さ $N$ の数列 $A = (A_1, A_2, ..., A_N)$ が与えられる。$1 \le i < j \le N$ かつ $A_i > A_j$ を満たす整数の組 $(i, j)$ の個数を求める
2. 単純な実装例
$N$ の制約が小さいとき、単純に2重ループで解くことができる。コードは次の通り。
N = int(input())
A = list(map(int, input().split()))
count = 0
for i in range(N):
for j in range(i + 1, N):
if A[i] > A[j]:
count += 1
print(count)
この実装では、すべてのペア $(i, j)$ を調べるため、計算量は $O(N^2)$ となる。
$N = 1000$ 程度なら問題ないが、$N = 10 ^ 5$ のような入力では、$10 ^ {10}$ 回の比較が必要になり、現実的な時間では終わらない。
そのため、より高速に転倒数を求めるアルゴリズムが必要となる。
3. BITによる高速化
Binary Indexed Tree(以下 BIT)とは、配列の累積和を動的に更新・取得できるデータ構造である。
通常の累積和は一度計算すると要素の変更に対応できないが、BITを使うと
- 要素の更新
- 区間和の計算
の両方を $O(\log N)$ で行うことができる。
この BIT を使って、転倒数を効率的に求めることができる。
考え方は次のとおりである。
長さ $N$ の数列 $A$ を右から順に見ていき、これまでに見た要素(=右側の要素)を BIT に追加していく。
各ステップで、自分より小さい要素が右側にいくつあるのかを BIT の累積和から求めることで、転倒数をカウントすることができる。
例として、配列 $A = [1, 5, 2, 4, 3]$ を考える。$A$ を右側から見ていき、各ステップで BIT の total(A[i] - 1) を計算することで、自分より小さい要素が右側にいくつあるかを求めることができる。
| 手順 | 見ている要素 ($=A[i]$) | BITに追加する前の状態 | 小さい要素の数 (total(A[i]-1)) |
転倒数の増分 | BIT更新後 |
|---|---|---|---|---|---|
| 1 | 3 | $[ \quad ]$ | 0 | +0 | $[3]$ |
| 2 | 4 | $[3]$ | 1(=3) | +1 | $[3,4]$ |
| 3 | 2 | $[3,4]$ | 0 | +0 | $[2,3,4]$ |
| 4 | 5 | $[2,3,4]$ | 3(=2,3,4) | +3 | $[2,3,4,5]$ |
| 5 | 1 | $[2,3,4,5]$ | 0 | +0 | $[1,2,3,4,5]$ |
4. BIT を使ったPythonによる実装
次の問題を解く。
Python による実装を次に示す。
class BIT:
def __init__(self, N):
self.N = N
self.bits = [0] * (self.N + 1)
def update(self, i, x):
while i <= self.N:
self.bits[i] += x
i += i & -i
def total(self, i):
res = 0
while i > 0:
res += self.bits[i]
i -= i & -i
return res
N = int(input())
A = list(map(int, input().split()))
bits = BIT(N)
count = 0
for i in range(N - 1, -1, -1):
count += bits.total(A[i] - 1)
bits.update(A[i], 1)
print(count)
5. 応用
今回の問題は、$1 \le N \le 150,000$ であり、配列 $A$ の値域も小さいため、BIT の配列をそのまま $O(N)$ の大きさで確保しても問題がなかった。($=$ 空間計算量が $O(N)$)
しかし、問題によっては $A_i$ の最大値が $10 ^ 9$ や $10 ^ 12$ のように非常に大きい場合があり、このようなケースで、値をそのまま BIT 配列の大きさで確保すると、メモリ不足になる。
そのような場合に用いられるのが、座標圧縮である。これは、値の大小関係を保ったまま、登場する値に連番の順位を付け直す手法である。
例として、配列 $A = [100, 10, 50, 10]$ を考える。
この配列に登場する値を昇順に並べると $[10, 50, 100]$ となる。それぞれに順位を割り当てると次のようになる。
| 元の値 | 新しい値(rank) |
|---|---|
| 10 | 1 |
| 50 | 2 |
| 100 | 3 |
これにより、配列 $A$ は $[3, 1, 2, 1]$ へと変換される。
この変換により、BIT のサイズは異なる値の個数だけ確保すればよくなる。
Python では次のように実装することができる。
comp = {v: i + 1 for i, v in enumerate(sorted(set(A)))}
A = [comp[a] for a in A]
このように、辞書を使って元の値と新しい値($=$ rank)を対応させる。
この手法が実際に必要となる代表的な問題として、次のようなものがある。
この問題では、制約から配列の取りうる値が $0 \le a_i \le 10 ^ 9$ となっているので、座標圧縮を行わなければ BIT を使用することができない。
Python では次のように実装することができる。
class BIT:
def __init__(self, N):
self.N = N
self.bits = [0] * (self.N + 1)
def update(self, i, x):
while i <= self.N:
self.bits[i] += x
i += i & -i
def total(self, i):
res = 0
while i > 0:
res += self.bits[i]
i -= i & -i
return res
N = int(input())
A = list(map(int, input().split()))
bit = BIT(N)
count = 0
# 座標圧縮
mapping = {a: i for i, a in enumerate(sorted(set(A)), start=1)}
# 転倒数のカウント
for i in range(N - 1, -1, -1):
# 座標圧縮後の値(rank)を用いてBITに反映
count += bit.total(mapping[A[i]] - 1)
bit.update(mapping[A[i]], 1)
print(count)
6. まとめ
今回は BIT を用いて転倒数を効率的に求める方法を解説した。さらに、応用として座標圧縮を併用する方法についても解説した。
転倒数の計測には、ここでは解説しなかったが、セグメント木(Segment Tree)を用いる方法も有効である。
いずれの場合も、単純な $O(N ^ 2)$ の全探索と比較して、大幅に高速化することができ、大規模データを扱う上で非常に有効である。