C++でinsertやeraseが高速な配列を作る

この記事は Aizu Advent Calendar 2016 2日目の記事です。

前の人は、@___deraさん、次の人は @hnjk さんです。

C++でよく使われる動的配列に、std::vectorがあります。std::vectorO(1)で配列の要素へのアクセス、push_backなどの操作が可能なデータ構造です。しかし、配列は要素がメモリ上で連続しているので、insertやeraseの処理はどうしても線形時間掛かってしまいます。 LinkedListであれば、リストの挿入先のポインタが分かればO(1)で挿入することが出来ます。しかし今度は反対に各要素へのアクセスが線形時間かかってしまいます。結局、任意の場所で挿入・削除をしようと思えば、これも線形時間かかってしまうことになります。

今回は、配列の要素へのアクセス速度を犠牲にO(logn)にする一方で、insertやeraseの処理をO(logn)まで高速化するデータ構造(競技プログラミングでよくあるテク)の解説をします。

参考文献は以下のスライドです。 プログラミングコンテストでのデータ構造 2 ~平衡二分探索木編~

erase()部分にバグがあるので、その部分の解説は後ほど追記します。

今回作るデータ構造とstd::vector, std::listとの主な違い

std::vector std::list 今回作るデータ構造
要素の挿入・削除の計算量 O(n) O(1) or O(n) O(log(n))
push_back(), (std::deque) push_front()の計算量 O(1) O(1) O(log(n))
要素へのアクセスの計算量 O(1) O(n) O(log(n))
メモリ上での要素の連続性

ご覧の通り、今回作るデータ構造はあらゆる操作がO(log(n))になっています。

平衡二分探索木を使う

平衡二分探索木というと、std::setにみられるように通常は値の集合の管理に使われます。しかし、既存のstd::setなどのデータ構造は利用できないので、平衡二分探索木を自作することになります。本記事では平衡二分探索木を動的配列のように利用する方法と、平衡二分探索木の具体的な実装方法について紹介します。折角なので、データ構造の名前はvectorとtreeを合わせてvectreeとでも名付けておきます。

赤黒木は書きたくない

有名な平衡二分探索木には、std::setでも使われている赤黒木や、AVL木などがあります。これらはを自作して拡張することで、目的の動的配列のようなデータ構造を作ることが出来ますが、実装が大変です。今回はより実装が軽量なTreapという平衡二分探索木を紹介します。

Treapの仕組み

Treapのノードには以下の2種類の情報があります。

  • 通常の二分探索木と同様のキーの値
  • 優先度を表す[0...1]の実数値

優先度の値には、そのノードが生成されると同時に[0...1]の範囲で一様なランダムの値が選ばれます。優先度が高いほど、そのノードが上に来るように木の回転を繰り返します。

f:id:moistx:20161201214318p:plain

このようにすると、なんと木の平衡が保たれて高さがO(logn)になります。

平衡化される理由について、Treapはランダムな順序で通常の二分探索木に要素を挿入してできる木と同様のものとみなすことができます。「ランダムな順序で構築された二分探索木」の高さがO(logn)となるのでTreapの高さもO(logn)といったように解釈することが出来ます。

基本的な実装

次に、基本的なvectreeの実装を見てみます。

namespace vectree
{

template<class T>
class vectree
{
  using treap    = detail::treap::node<T>;
  using node_ptr = typename treap::node_ptr;

  node_ptr root {nullptr};

public:

  vectree() = default;

  ~vectree()
  {
    if (root) { delete root; }
  }

  auto size() const noexcept
  { return root ? root->size() : 0; }

  T operator[] (unsigned position) const { return treap::nth_element(root, position)->value(); }

  auto insert(unsigned position, T val)
    -> void
  {
    root = treap::insert(root, position, val);
  }

  auto push_back(T val)
    -> void
  {
    root = treap::insert(root, size(), val);
  }

  auto push_front(T val)
    -> void
  {
    root = treap::insert(root, 0, val);
  }

  auto erase(unsigned position)
    -> void
  {
    root = treap::erase(root, position, nullptr);
  }

  auto clear()
    -> void
  {
    if (root) delete root;
    root = nullptr;
  }
};

}

vectreeのの実体は、Treapの根のノードを1つ確保して、vectreeでinsert(), erase()したらTreapでもinsert(), erase()するということを書いているだけです。operator[]()で呼び出しているnth_element()で、配列の指定した添字の要素にアクセスしています。 動的配列に明らかに必要なIteratorの実装が書いてありませんが、insert()nth_element()などが書ければ多分書けると思うので省略しました。

次にTreapの基本的な実装について見てみます。

Treapのノードは以下の4種類の情報を持ちます。

  • ノードの値 val
  • 左の子、右の子
  • 優先度
  • 自分を頂点とした部分木のサイズ

部分木のサイズは、insert()nth_element()などの処理で必要になります。木の回転やノードの挿入・削除が起こるたびに再計算をします。

基本的なtreapのヘッダ内容は以下のとおりです。

namespace vectree::detail::treap
{

template<class T>
class node
{
  T val;
  node<T> *lch, *rch;
  double priority;
  unsigned tree_size;

  ...
  explicit node(T v, double p)          noexcept
  : val{v}, lch{nullptr}, rch{nullptr}, priority{p}, tree_size{1}
  {}

  static auto erase_node(node<T>* deletee) -> node<T>*;

public:

  using node_ptr = node<T>*;
  using node_ptr_reference = node_ptr&;
  using node_ptr_const_reference = node_ptr const&;

  ~node()
  {
    if (lch) delete lch;
    if (rch) delete rch;
  }

  auto& value()               noexcept
  { return val; }

  auto const& value()   const noexcept
  { return val; }

  auto size()           const noexcept
  { return tree_size; }

  auto expired_left()   const
  { return !lch; }

  auto expired_right()  const
  { return !rch; }

  auto expired_both()   const
  { return expired_left() && expired_right(); }

  auto left()               -> node_ptr_reference;
  auto right()              -> node_ptr_reference;
  auto child(side w)        -> node_ptr_reference;

  static auto update_size(node_ptr t)                         -> node_ptr;
  static auto rotate(node_ptr t, side right)                  -> node_ptr;
  static auto insert(node_ptr tree, unsigned position, T val) -> node_ptr;
  static auto erase(node_ptr tree, unsigned position)         -> bool;
  static auto nth_element(node_ptr tree, unsigned position)   -> node_ptr;

  static auto create(T val)
  {
    static auto rand_priority_gen = std::bind(
      std::uniform_real_distribution<>{0.0, 1.0},
      std::default_random_engine{});
    return new node(val, rand_priority_gen());
  }

};
}

ヘッダの内容がわかったところでinsert()erase()nth_element()を実装していきます。

ただ、その前に肝心の平衡二分探索木をいかにして配列として利用するかについて説明する必要があります。

平衡二分探索木をどう使えば配列として管理できるか

平衡二分探索木は、平衡な木の上で二分探索をするから、高速に要素の検索・挿入・削除が出来ます。しかし、配列に入れる要素は、普通はソートされていません。ソートされていなければ、二分探索木上に配置することが出来ません。 配列の要素はソートされていませんが、どんな値が要素となっても、配列には"ソート"されているものがあります。それは配列の添字です。配列の添字は、確かに左から順に0, 1, 2, ... と辞書順に増加していくので、これは探索に使えそうです。

配列の添字を二分探索のキーとすれば添字の要素の特定できるか?

以下の図のように、要素を添字と値のペアだとみなして、添字の順序で木の上を探索すれば良さそうな気がします。

f:id:moistx:20161202004340p:plain

確かに、この木の上では二分探索が可能です。しかし、例えば添字1と2の間に新しい要素の7を挿入しようとした場合、添字はいくつとすべきでしょうか。 添字が2以降の要素に対し、添字の値を全て1ずつ増加させるのは本末転倒です。となると、以下のように添字1.5なるもの導入すればよいでしょうか。

f:id:moistx:20161202004339p:plain

一見良さそうに見えますが、この方法ではうまくいきません。コンピュータ上で限りなく細かい実数値を扱うことは出来ないので、何度も挿入を繰り返すと簡単に添字に誤差が出てしまいます。

添字の要素の特定には、木のサイズを活用する

実は、ノードに添字の情報を持たせる必要はありません。木のサイズを活用することで、添字を特定することが出来ます。

ノードに中間順巡回で番号を振り、その上で通常の二分探索木の検索と似た方法で、指定した添字の要素を検索することが出来ます。

例えば、以下の図で5番のノード(添字5の要素)を特定するとします。

f:id:moistx:20161201231134p:plain

根について木のサイズは12で、根の左部分木に6個、右部分木に5個、根の頂点それ自身が1個です。根から探索を開始します。はじめの6個のノードは左部分木にあるので、5番のノードは左部分木に存在するはずです。よって、左部分木に遷移します。

今度は、左部分木に3個、右部分木に2個のノードがあります。5番目のノードは左部分木のサイズより大きいので、左部分木でもなく現在の場所でもなく、右部分木の中に存在するはずです。(左部分木のサイズ) + 1 = 4 つのノードを探索する必要が無くなりました。よって、右部分木の頂点から見て 5 - (3 + 1) = 1番目目のノードを特定すれば良いことになります。

f:id:moistx:20161201230926p:plain

右部分木に遷移します。現在の位置から見て左部分木にノードが1個、右部分木にはノードがなく、現在の場所に1個のノードが存在します。左部分木のサイズが1で、1番目のノードを特定したいので、現在の場所が特定したいノードとなります。

以上のように、根から辿って特定したいノードの番号(添字の要素)を右部分木に遷移するときだけ (左部分木のサイズ) + 1だけ減算していけば、指定したノードの番号、つまり指定した添字を持つ要素がO(木の高さ) = O(logn)で特定できることとなります。

以上より、添字の要素のノードを特定するnode<T>::nth_element()を実装してみましょう。 まず左部分木のサイズと今の部分木の何番目かを示すpositionとを比較し、左右どちらの部分木に進行するかを決めます。次に、左に進行する場合はpositionを変えずに進行し、右に進行する場合は左の部分木と今いるの頂点の数を引いた分だけ残して進行しています。

template<class T>
auto node<T>::nth_element(node<T>::node_ptr tree, unsigned position)
  -> node<T>::node_ptr
{
  if(!tree)
  {
    return nullptr;
  }

  auto const left_size = tree->expired_left() ? 0 : tree->left()->size();

  if(left_size == position)
  {
    return tree;
  }

  auto const to_right = left_size < position;
  auto const subtree_pos = position - (to_right ? 1 + left_size : 0);
  
  return nth_element(to_right ? tree->right() : tree->left(), subtree_pos);
}

node<T>::insert()も実装してみましょう。 nth_element()の処理に、新しいノードの追加と、そのノードを優先度が適切な位置になるまで、木の回転で上方に移動させていく処理が加わりました。

また、部分木に新しいノードが挿入された場合、木のサイズが変わるので、update_size()を呼び出しています。

template<class T>
auto node<T>::insert(node<T>::node_ptr tree, unsigned position, T val)
  -> node<T>::node_ptr
{
  if (!tree)
  {
    return tree = node<T>::create(val);
  }

  auto const left_size = tree->expired_left() ? 0 : tree->left()->size();
  auto const to_right = left_size < position;

  auto const next_tree = to_right ? tree->right() : tree->left();
  auto const subtree_pos = position - (to_right ? 1 + left_size : 0);
  auto const inserted_subtree = insert(next_tree, subtree_pos, val);

  (to_right ? tree->right() : tree->left()) = inserted_subtree;  // 各ノードが親のポインタを持っていないので、張替えが必要

  update_size(tree);

  if (tree->priority < inserted_subtree->priority)
  {
    return rotate(tree, to_right ? side::left : side::right);
  }

  return tree;
}

update_size()の実装は単純で、(左部分木のサイズ) + (右部分木のサイズ) + 1 で更新すれば良いです。

template<class T>
auto node<T>::update_size(node<T>::node_ptr t)
  -> node<T>::node_ptr
{
  t->tree_size =
    1 + (t->expired_left()  ? 0 : t->left()->size())
      + (t->expired_right() ? 0 : t->right()->size());
  return t;
}

次に、木の回転rotate()の実装です。引数dirは左回転か右回転かを表します。左回転する場合は、ノードtreeから見て右の子との位置関係を関係を変えることになるので、node<T>::sibling()という関数でdir = side::leftならside::rightを返すようにしています。右回転する場合も同様です。

template<class T>
auto node<T>::rotate(node<T>::node_ptr tree, side dir)
  -> node<T>::node_ptr
{
  auto const subtree = tree->child(node<T>::sibling(dir));
  tree->child(node<T>::sibling(dir)) = subtree->child(dir);
  subtree->child(dir) = tree;
  update_size(tree);
  update_size(subtree);
  return subtree;  // 今回の実装では各ノードが親ポインタを持っていないので、処理の後に回転後に上に来たノードを返してやる実装になっている
}

右回転の場合、以下のように辺の張替えが行われます。yがrotate()でいうところのtreeで、xが新たに返却されるノードとなります。 f:id:moistx:20161202051229p:plainf:id:moistx:20161202051230p:plainf:id:moistx:20161202051232p:plain

最後にerase()を実装してみましょう。これもinsert()同様にnth_element()を拡張させることが基本になります。

※この部分はバグがあったため後ほど追記します。

テスト・デバッグ

一通りの作業が終わったので、簡単にテストコードを書いてみましょう。 木がいい感じに平衡化されているかを確かめたいので、木を標準出力する関数を書いてみます。

template<class T>
inline void debug_recur(node<T>* t, int depth = 0) {

  for(int i=0; i<depth * 2; i++)
    std::cout << " ";

  if(!t)
  {
    std::cout << "(nil)" << std::endl;
    return;
  }

  std::cout << "(" << t->value() << "[" << t->pri() << "]" << std::endl;  // 優先度も正しく表示されているか確認したいので、デバッグ用にpri()を用意した

  debug_recur(t->left(), depth + 1);
  debug_recur(t->right(), depth + 1);

  for(int i=0; i<depth * 2; i++)
    std::cout << " ";

  std::cout << ")" << std::endl;
}

[0, 1, 2,...] と増加する値を順に挿入して表示してみると、そこそこ平衡化されていることが確認できると思います。Treapはランダムに二分探索木に挿入したものと同じなので、厳密に平衡化されているわけではないですがより多くのノードを挿入しても深さは余り増えません。また、優先度も上に来るノードほど高く、下に来るノードほど小さくなっていることが確認出来ます。

入力
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]

出力 (だいたい平衡化できてる)
(8[0.995085]
  (1[0.891611]
    (0[0.0850324]
      (nil)
      (nil)
    )
    ...
    )
    (nil)
  )
)

次に、vectreeとstd::vectorとで簡単に速度比較をしてみます(eraseはバグが取れたら後ほど追記します。すみません)。

方法は、各々のデータ構造に105個の要素をpush_back()し、次に各々のデータ構造の中央に106個の要素をinsert()して確かめます。 vectreeのpush_back()が少し遅く、insertがとても速くなれば期待通りです。

std::vector vectree
空の配列に105個のpush_back() 0msec 28msec
同じ配列に中央に106個のinsert() 57203msec 472msec

期待通りの結果が得られました。

使用したコードは以下のとおりです。

  // push_back(), push_front(), insert() のテスト
  vectree::vectree<int> vs;

  int const FirstSize  = 100000;
  int const SecondSize = 1000000;

  // vectree
  std::cout << "---- vectree ----" << "\n";
  std::cout << FirstSize << "個の要素をpush_back()" << "\n";
  auto start = std::chrono::system_clock::now();
  for(int i=0; i<FirstSize; i++)
    vs.push_back(i);
  auto end = std::chrono::system_clock::now();
  std::cout << "Elasped time(msec) = " << std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count() << "\n";

  std::cout << "加えて、配列の中央あたりに" << SecondSize << "個の要素をinsert()" << "\n";
  start = std::chrono::system_clock::now();
  for(int i=0; i<SecondSize; i++)
    vs.insert(vs.size() / 2, i);
  end = std::chrono::system_clock::now();
  std::cout << "Elasped time(msec) = " << std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count() << "\n";

  // std::vector
  std::vector<int> vec;
  std::cout << "--- std::vector ---" << "\n";
  std::cout << FirstSize << "個の要素をpush_back()" << "\n";
  start = std::chrono::system_clock::now();
  for(int i=0; i<FirstSize; i++)
    vec.push_back(i);
  end = std::chrono::system_clock::now();
  std::cout << "Elasped time(msec) = " << std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count() << "\n";

  std::cout << "加えて、配列の中央あたりに" << SecondSize << "個の要素をinsert()" << "\n";
  start = std::chrono::system_clock::now();
  for(int i=0; i<SecondSize; i++)
    vec.insert(vec.begin() + vec.size() / 2, i);
  end = std::chrono::system_clock::now();
  std::cout << "Elasped time(msec) = " << std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count() << "\n";

  std::cout << "配列に入っている値が同じか: ";
  for (int i=0; i<vec.size(); i++)
  {
    assert(vs[i] == vec[i]);
  }

  std::cout << "OK" << std::endl;

おわりに

長々と書いたものの、結局は参考文献のスライドの一部についての説明に対応しています。そちらを読んだほうが分かりやすいかもしれません。参考にしたスライドはinsert()erase()に加えて、merge(), split()という操作を可能にさせています。配列を連結・分割する操作をO(logn)で出来るようにしたものです。興味がある方はそちらを含めて読んでみると良いかもしれません。