マージソート木の概要と実装
マージソート木とは、長さ $N$ の静的な数列 $A$ に対して以下の操作を $O(\log^2 N)$ で処理するデータ構造です。
- 区間 $[l,r)$ において、 $A_i \leq x$ を満たす要素数を求める
- 区間 $[l,r)$ において、 $a \leq A_i < b$ を満たす要素数を求める
マージソート木はセグメント木と同じように区間を分割して、各分割ブロックに該当する区間の数列をソートした状態で保持します。
マージソートの要領で構築できるので、構築は $O(N \log N)$ 時間で行えます。
数列をソートした状態で持つので空間計算量がやばそうに思えますが、木の深さは $O(\log N)$ なので空間計算量も $O(N \log N)$ となります。
各クエリに対しては、セグメント木と同様に区間を包含するノードを探し、その区間のソートされた数列に対して二分探索を行うことで解答できます。
区間を包含するノードの探索に $O(\log N)$、各探索の二分探索に $O(\log N)$ 時間かかるので、各クエリに対して $O(\log N \times \ \log N) = O(\log^2 N)$ 時間で解答できます。
$a \leq A_i < b$ を満たす要素数を計算するためには毎回 2 回の二分探索が必要になるため、定数倍が悪化します。
Pythonの実装例
from bisect import bisect_right,bisect_left
class MergeSortTree:
def __init__(self, data):
self.n = len(data)
self.size = 1 << self.n.bit_length()
self.tree = [[] for _ in range(2 * self.size)]
for i in range(self.n):
self.tree[self.size + i] = [data[i]]
# 子ノードの要素をソートしてすべて持つように構築
for i in range(self.size - 1, 0, -1):
# マージソートの要領でマージ
left_child = self.tree[2 * i]
right_child = self.tree[2 * i + 1]
l,r = 0,0
while l < len(left_child) and r < len(right_child):
if left_child[l] < right_child[r]:
self.tree[i].append(left_child[l])
l += 1
else:
self.tree[i].append(right_child[r])
r += 1
while l < len(left_child):
self.tree[i].append(left_child[l])
l += 1
while r < len(right_child):
self.tree[i].append(right_child[r])
r += 1
# 区間 [l, r) で値が x 以下の要素数を求める
def query_leq(self, l, r, x):
l += self.size
r += self.size
res = 0
while l < r:
if (l&1):
res += bisect_right(self.tree[l], x)
l += 1
if (r&1):
res += bisect_right(self.tree[r-1], x)
r -= 1
l >>= 1
r >>= 1
return res
# 区間[l, r) で値が a 以上 b 未満の要素数を求める
def query_range(self, l, r, a, b):
l += self.size
r += self.size
res = 0
while l < r:
if (l&1):
res += bisect_left(self.tree[l], b) - bisect_left(self.tree[l], a)
l += 1
if (r&1):
res += bisect_left(self.tree[r-1], b) - bisect_left(self.tree[r-1], a)
r -= 1
l >>= 1
r >>= 1
return res
C++の実装例
struct MergeSortTree {
int n;
int size;
vector<vector<ll>> tree;
MergeSortTree(const vector<ll>& data) {
n = data.size();
size = 1;
while(size < n) size *= 2;
tree.assign(2 * size, vector<ll>());
// 葉に値をセット
for (int i = 0; i < n; i++) {
tree[size+i] = {data[i]};
}
// 残りは空の状態で初期化しておく
for (int i = n; i < size; i++) {
tree[size+i] = vector<ll>(); // 空
}
// 内部ノードをマージソートの要領で構築
for (int i = size - 1; i >= 1; i--) {
const vector<ll>& left_child = tree[2*i];
const vector<ll>& right_child = tree[2*i+1];
tree[i].resize(left_child.size() + right_child.size());
int idx = 0;
int l = 0, r = 0;
while(l < (int)left_child.size() && r < (int)right_child.size()){
if(left_child[l] < right_child[r]){
tree[i][idx++] = left_child[l++];
} else {
tree[i][idx++] = right_child[r++];
}
}
while(l < (int)left_child.size()){
tree[i][idx++] = left_child[l++];
}
while(r < (int)right_child.size()){
tree[i][idx++] = right_child[r++];
}
}
}
// 区間 [l, r) において、値が x 以下の要素数を返す
int query_leq(int l, int r, ll x) {
int res = 0;
l += size;
r += size;
while(l < r){
if(l & 1) {
res += upper_bound(tree[l].begin(), tree[l].end(), x) - tree[l].begin();
l++;
}
if(r & 1) {
r--;
res += upper_bound(tree[r].begin(), tree[r].end(), x) - tree[r].begin();
}
l /= 2;
r /= 2;
}
return res;
}
// 区間 [l, r) において、値が [a, b) の要素数を返す
int query_range(int l, int r, ll a, ll b) {
int res = 0;
l += size;
r += size;
while(l < r){
if(l & 1) {
res += (lower_bound(tree[l].begin(), tree[l].end(), b) -
lower_bound(tree[l].begin(), tree[l].end(), a));
l++;
}
if(r & 1) {
r--;
res += (lower_bound(tree[r].begin(), tree[r].end(), b) -
lower_bound(tree[r].begin(), tree[r].end(), a));
}
l /= 2;
r /= 2;
}
return res;
}
};
使える問題(ほかにもあったと思うので見つけたら追記します)
動的マージソート木について
残念ながら 2024 年 11 月現在、競技プログラミングにおいて動的マージソート木が必要になる問題はまだないようで、動的マージソート木について言及された記事もほとんどありません。12
マージソート木は静的な配列を扱うものですが、要素変更可能にすることもできると思います。
ソートされた数列を tatyam set などで持つことで要素の変更を $O(\log^2 N)$ などでおこなうことができ、クエリの解答も変わらず $O(\log^2 N)$ で可能だと思います。
たぶんAIに投げたら実装してくれる気がします。