木構造の実装テクニック

この記事は Competitive Programming Advent Calendar 2018 の 2 日目の記事として書かれました。
ところで今日は 26 日ですね。大変遅れて申し訳ありませんでした。

はじめに

タイトルの主語がでかいですが競技プログラミングで C++ の生ポインタを使って実装する際の話です。

木構造のデータ構造は競技プログラミングでたまに必要になりますが、実装が複雑になりがちでバグりやすく、また平衡二分探索木などの重めの木だと想定解法と同じ計算量でも TLE する事がよくあります。この記事では、そのようなイライラを回避するための実装を楽にするテクニックやチューニングテクニックをいくつか紹介します。

木構造と言っても、この記事ではポインタを使った実装を前提に話を進めます。したがって、平衡二分探索木などを実装するときには参考になるかもしれませんが、二分ヒープやセグメントツリーのように配列のインデックス整数を器用に使うようなものは対象外です。

tl; dr

一番下のベンチマークのところにコンパイルが通るコード全体を貼ったので忙しい人はそちらをどうぞ。

秋葉先生による平衡二分探索木のバイブル的資料

この記事では平衡二分探索木の Treap を扱います。それについて知らない人は次のスライドを見てください。以降このスライドを読んでる前提で話が進みます。

プログラミングコンテストでのデータ構造 2　～平衡二分探索木編～ Takuya Akiba

ノード構造体の定義

この記事では、RMQ (Range Mimimum Query) を処理する Treap を説明に使います。
次のようなノード構造体の定義を使うものとします。

using random_type = uint_fast32_t;
using key_type = int;
const key_type INF = 1000000000;
const int MAX_N = 1000010; // 挿入の回数 + 少し

struct Node {
    key_type key;
    random_type priority;
    Node *left, *right;
    int size;
    key_type min;

    Node() {}

    Node(const key_type &x) : Node(x, xor128(), nil, nil, 1, x) {} // nil については後述 (ヌルポみたいなものだと思ってください)

    Node(const key_type &key_, random_type priority_, Node *left_, Node *right_,
         int size_, const key_type &min_)
        : key(key_), priority(priority_), left(left_), right(right_), size(size_),
          min(min_) {}
};

using np = Node *;
using npair = std::pair<np, np>;

ポインタ型に別名をつける

C++ の有名な罠の一つに次のようなものがあります。

// a はポインタ型だが b は Node 型なのでコンパイルエラー！
Node *a = new Node(0), b = new Node(1);

多くの場合はコンパイル時に気づけますが、暗黙の型変換などのせいで、稀に気づけないケースもあります。そのようなミスを防ぐため、またタイプ数を減らすために、次のような別名を付けましょう。

using np = Node *;
// a, b ともにポインタ型
np a = new Node(0), b = new Node(1);

split の戻り値など、ポインタのペアを扱うことも多いので pair<np, np> にも名前をつけましょう。

using npair = std::pair<np, np>;

関数の仕様を意識する

木の実装に限ったことではありませんが、再帰関数を書くときは関数の仕様を最初にきちんと定義しましょう。
実装する際は、何を引数にとって何を返し、戻り値は何を保証するのか、どんな副作用があるのかを強く意識するとバグを生みづらくなります。

例えば、冒頭に貼った秋葉さんのスライドにおいて、merge と split などという関数は以下のことを意識すると読みやすいです。

• 引数は正しい
  • そうでないものは渡してはいけない
• 内部では一時的に正しくなくてもよい
• return するノードは正しい

ここで、「正しい」とは要素数や最小値の値のフィールドが正しい、左右のバランスが取れているといったことを指します。

再帰呼び出しを使ったコードを追うのは慣れないとややこしく感じるかもしれませんが、定義だけを意識して自身を呼んでいる先を深読みしないのがコツです。

また、バグったら次のようなチェック用関数を書いて怪しい箇所に割り込ませるのがいいです。
Node::nil の説明がまだですが nullptr と似たようなものです。

#include <cassert>
void verify(np n) {
    if (n == Node::nil)
        return;
    assert(n->size == 1 + n->left->size + n->right->size);
    assert(n->min == std::min({n->left->min, n->key, n->right->min}));
    assert(n->left != nullptr);
    assert(n->right != nullptr);
    assert(n->priority >= n->left->priority);
    verify(n->left);
    assert(n->priority >= n->right->priority);
    verify(n->right);
}

void buggy() {
    // (バグってそうなコード)
    #ifdef DEBUG
    verify(node);
    #endif
}

O(n) 時間で構築

扱う数列が最初に与えられるとき、次のように毎回末尾に挿入すると $O(N \log N)$ 時間かかります。

int N; cin >> N;
np root = nullptr;
for(int i = 0; i < N; ++i) {
    int x;
    cin >> x;
    root = insert(root, i, x);
}

一方、葉から順に構築すると $O(N)$ でできます。

np make_tree(key_type *left, key_type *right) {
    int sz = right - left;
    if (sz == 0)
        return Node::nil;
    key_type *mid = left + sz / 2;
    np lc = make_tree(left, mid);
    np rc = make_tree(mid + 1, right);
    return new Node(*mid, -1, lc, rc, sz, std::min({lc->min, *mid, rc->min}));
}

int x[100010];
int main() {
    int N; cin >> N;
    for(int i = 0; i < N; ++i) {
        cin >> x[i];
    }
    np tree = make_tree(x, x + n);
}

RBST なら簡単で上のコードで終わりですが、Treap の場合はヒープの条件を満たさせる必要があります。

// make_tree の定義は上と同じ
np make_treap(key_type *left, key_type *right) {
    np root = make_tree(left, right);
    int n = right - left;
    std::vector<random_type> ps(n);
    std::generate(ps.begin(), ps.end(), xor128);
    std::sort(ps.begin(), ps.end());
    std::queue<np> que;
    que.push(root);
    while (que.size()) {
        np n = que.front();
        que.pop();
        int pri = ps.back();
        ps.pop_back();
        n->priority = pri;
        if (n->left != Node::nil) que.push(n->left);
        if (n->right != Node::nil) que.push(n->right);
    }
    return root;
}

ソートに $O(N \log N)$ かかっているだろうという突っ込みが想定されますが insert よりは速いです。
シンプルに線形構築する方法はあるんでしょうか。僕は思いつかなかったので、思いついた方がいれば教えてください (どうしても「乱数」をソートする必要があるので無理な気がしますが)。

メモリプール

少量のメモリを頻繁に確保/解放するのは効率が悪いです。あらかじめ大きな領域を貯めて (プールして) おき、そこから少しずつ使い、解放も最後にまとめて行いましょう。

Node 構造体の定義に次のコードを追加します。operator new のオーバーライドという機能を使っているので、詳しく知りたい方は「operator new override c++」などで検索してください。

struct Node {
    // (略)
    
    static int node_count;
    // 静的に確保した配列から返す
    void *operator new(std::size_t) {
        static Node pool[MAX_N];
        return pool + node_count++;
    }
};

int Node::node_count = 0;

メモリプールを使ったときは new した個々のノードは delete すると segmentation fault が起こるので放置しましょう。
一括解放なら node_count を 0 に戻すことでできます。

struct Node {
    // (略)
    static void delete_all() { Node::node_count = 0; }
};

nullptr の代わりにダミーノードを使って update の分岐を無くす

頂点に子が無いことを nullptr で表現することが多いと思います。そのとき、あるノード以下の部分木の要素数を求める関数 count や最小値を求める関数 min は次のように書かれることが多いです。

int count(np n) {
    return n == nullptr ? 0 : n->cnt;
}

key_type min(np n) {
    return n == nullptr ? INF : n->min;
}

np update(np n) {
    n->size = 1 + size(n->left) + size(n->right);
    n->min = std::min({min(n->left), n->key, min(n->right)});
    return n;
}

ここで何も指さないことを表すダミーのインスタンス nil を用意すると、上の関数 count と min は不要になり、三項演算子の分岐がなくなるので少し速くなります。

struct Node {
    // (略)
    static Node *const nil;
};
Node *const Node::nil = new Node(key_type(), std::numeric_limits<random_type>::min(),
                                 nullptr, nullptr, 0, INF);

np update(np n) {
    n->size = 1 + n->left->size + n->right->size;
    n->min = std::min({n->left->min, n->key, n->right->min});
    return n;
}

メモリプールと併用するときは、全解放のときに nil を解放しないように注意してください。
つまり、nil インスタンスの確保を最初に行い node_count = 0 だったところを node_count = 1 に変えてください。

lower_bound と upper_bound

ソートされた列と引数 x に対して、lower_bound は x 以上の最小の要素を返す関数で、upper_bound は x より大きい最小の要素を返す関数です。
C++ 使いの人にはおなじみだと思います。

ネット上にアルゴリズムが解説されている記事が少ないですが、以下のようにすればできます。ただし、木 n が管理する列はソートされていることを仮定します。

np lower_bound(np n, const key_type &x) {
    if (n == Node::nil)
        return Node::nil;
    else if (n->key >= x) {
        np res = lower_bound(n->left, x);
        return res != Node::nil ? res : n;
    } else
        return lower_bound(n->right, x);
}

np upper_bound(np n, const key_type &x) {
    if (n == Node::nil)
        return Node::nil;
    else if (n->key > x) {
	np res = upper_bound(n->left, x);
        return res != Node::nil ? res : n;
    } else
        return upper_bound(n->right, x);
}

2 つの違いは 3 行目の else if の中の不等号のみです。

永続 RBST は偏る

らしいです。赤黒木や AVL 木を使いましょう。

RBSTがコピー可能は嘘という疑惑

乱数アルゴリズム

速い乱数アルゴリズムを使いましょう。私は軽実装かつ速い XorShift を使っています (他のをよく知らないので)。

using random_type = uint_fast32_t;
random_type xor128() {
    static random_type x = 123456789, y = 362436069, z = 521288629, w = time(0); // seed は w に入れる
    random_type t = x ^ (x << 11);
    x = y;
    y = z;
    z = w;
    w = (w ^ (w >> 19)) ^ (t ^ (t >> 8));
    return w;
}

チャレンジでハックされないくらい乱雑ならわりとなんでもいいです。

標準ライブラリの rand はダメです。線形合同法が用いられていることが多く、剰余演算を含むので若干遅いです。

さらに、VisualStudio や Codeforces の一部言語設定だと RAND_MAX が 65535 です (私は VS ユーザーではなく 2018 年に人に聞いた話なので、今実際にどうかは各自確認していただきたいです)。すなわち 65536 以上の値が得られないので、$10^5$ オーダーのノード数の Treap では容易に衝突します。

このように、何かと問題が多いので最新の C++ では非推奨となっています。なので、競技プログラミングなら上で紹介したような自前実装で十分ですし、標準ライブラリを使いたい場合は std::mt19937 などの新しい乱数ライブラリが推奨されています。

とは言っても所詮競技プログラミングなので、慣れているオンラインジャッジで rand でも動くと分かっているならば使ってもいいと思います。

定数に const/constexpr を付ける

付けられるときになんとなく付けておくと、コンパイラが不変なことを利用した最適化をしてくれることがあります。

const key_t INF = 1000000000
const int MAX_N = 1000010;
const int MOD = 1000000007;

関数に inline を付ける

関数に inline を付けると速くなることがあります。

inline np update(np n) { /* (略) */ }
inline random_type xor128() { /* (略) */ }

まともなオンラインジャッジでは最適化オプションが付いているので効果が無いことが多いです。POJ では効果が顕著です。

ベンチマーク

冒頭のスライドのシンプルな Treap と上で紹介した高速化手法込みの Treap を比較してみます。スライドのコードは入門者用の究極にシンプルさを意識した実装なので、そもそも比べるのは不公平だということは理解してください。

条件

次の問題を解く時間を測ります。

長さ $N$ の数列 $a$ が与えられます。数列に対する $Q$ 個のクエリを順に処理してください。
クエリには次のような type0, type1 の $2$ 種類あり、全ての type0 クエリの結果の XOR を出力してください。

• type0: 区間 $[l, r)$ の最小値を求める (RMQ)
• type1: $a$ の前の $k$ 個の要素を後ろに移動する (rotate)
  $N,Q=10^6$, $0 \leq a_i \leq 10^9$ を満たします。

例えば $a=[1,2,3,4,5]$ に対して $k=2$ の type1 クエリを適用すると $a=[3,4,5,1,2]$ となります。

この問題の解答の関数 merge, split, range_min (RMQ) の実行時間 (入出力と Treap の構築を除いた時間) を測ります。
もっとちゃんとやるべきですが省略します。
ジェネレータを含むコードはこちらです。ベンチマークについて詳しくはこれを参照してください (雑でごめん)。
https://gist.github.com/tubo28/bf0ea645525c8b75fd82f3e7568a7341

• basic.cc : ほぼ冒頭のスライドを写経したもの
• fast.cc : このスライドで紹介した高速化手法込みのもの
• generate.rb : テストケースジェネレータ

環境は次の通りです。

• OS: Ubuntu 18.04 (VMware Workstation 14 Player)
• CPU: AMD Ryzen 7 1800X Eight-Core Processor
• RAM: 2GB
• コンパイラ: gcc (Ubuntu 7.3.0-27ubuntu1~18.04) 7.3.0
• コンパイルオプション: g++ -std=c++14 -O3 ${src}.cc

結果

$ ruby generate.rb > in
$ g++ -std=c++14 -O3 basic.cc && ./a.out < in
42991693
1661 ms
$ g++ -std=c++14 -O3 fast.cc && ./a.out < in
42991693
930 ms

55% くらいの実行時間になりました。

おわりに

色々と紹介してきましたが、このようなチューニングはコンテスト中にやるのは無駄な努力に終わることが多いのであまりよくありません。より本質的な高速化ができないか再考しましょう。
しかし、用意があるといざというときに有利なのでライブラリ化しておきましょう。盆栽の手入れは楽しいです。

以上です。それでは良いお年をお迎えください。