wat-arrayを使った2次元探索プログラム

岡野原氏の作成したwavelet木を使った高速配列処理ライブラリwat-arrayを利用して、2次元探索のプログラムを書いてみた。
なお、自分はwavelet木のアルゴリズムについては全く分かってないですが、wat-arrayでは配列に対して、操作を行うインターフェイスがしっかり与えられているのでそれを見ながら作りました。

問題定義

2次元座標の集合P={(x,y)}が与えられる。Queryとして(xs,xe,ys,ye)が与えられたときにPの中でxs <= x < x, ys <= y < yeを満たす点の数を答えるというものを考える。

なお、Pの内容は途中で変化したりすることはないものとする(変更が加わった場合は一から作り直す)。

インターフェースとしては次のようになる、2次元座標の表現にはpairを用いることにする。

namespace wat2DSearch{
  typedef std::pair<uint64_t , uint64_t> Point;
  class Wat2DSearch{
    public:
      Wat2DSearch(const std::vector<Point> & P);
      //Compute the frequency of points where xs <= x < xe and ys <= y < ye
      uint64_t FreqRange2D(uint64_t xs , uint64_t xe , uint64_t ys , uint64_t ye);
  };
}

使い方の例

int main(){
  vector<wat2DSearch::Point> P;
  P.push_back(make_pair(1,2));
  P.push_back(make_pair(2,1));
  P.push_back(make_pair(2,4));
  P.push_back(make_pair(4,2));
  wat2DSearch::Wat2DSearch wat(P);
  cout << wat.FreqRange2D(0 , 5 , 0 , 5) << endl;
  cout << wat.FreqRange2D(2 , 3 , 2 , 4) << endl;
  cout << wat.FreqRange2D(2 , 3 , 2 , 5) << endl;
  return 0;
}

以降でどのようにしてこのFreqRange2Dを実現するかについて述べる。

アルゴリズム

Pの値をxの昇順でソートする。たとえばP={(2,1),(4,2),(1,3),(2,4)}のときソート後は

index	0	1	2	3
座標	(1,3)	(2,1)	(2,3)	(4,2)

となる。これをx座標とy座標に分解すると

index	0	1	2	3
x座標	1	2	2	4
y座標	3	1	4	2

となる。

たとえばxs=2,xe=3のとき、x座標が昇順に並んでいることからy座標の配列のうちindexが[１ー３）までの範囲に関してys <= y < yeを満たすものの数を求めることと同じであることが分かる。
与えられた部分配列の中で指定されたRangeに収まるような要素の数を求めるのはwatArray::FreqRangeを用いて求めることができる。
また、与えられたxs,xeに対してindexを求める部分もxが昇順に並んでいることに着目するとxsに対応するindexはx座標の配列の中でxsより小さいものの数であることが分かり、
これはwatArray::FreqSumを用いて簡単に求まる。xeについても同様である。

プログラム

以上を踏まえて作ったのが以下のコードである、watArrayを使うことにより簡潔に書けていることが分かる。

ヘッダ

#include <wat_array/wat_array.hpp>
#include <utility>
#include <vector>
namespace wat2DSearch{
  typedef std::pair<uint64_t , uint64_t> Point;
  class Wat2DSearch{
    public:
      Wat2DSearch(const std::vector<Point> & P);
      uint64_t FreqRange2D(uint64_t xs , uint64_t xe , uint64_t ys , uint64_t ye);
    private:
      wat_array::WatArray xwat;
      wat_array::WatArray ywat;
  };
}

実装部分(2011/1/11 FreqRangeの境界条件での仕様に対応)

#include "Wat2DSearch.hpp"
#include <algorithm>
using namespace std;
namespace wat2DSearch{
  Wat2DSearch::Wat2DSearch(const vector<Point> & P): xwat(), ywat(){
    vector<Point> tmp(P);
    sort(tmp.begin() , tmp.end());
    vector<uint64_t> xArray , yArray;
    for(vector<Point>::iterator it = tmp.begin() ; it != tmp.end() ; ++it){
      xArray.push_back(it->first);
      yArray.push_back(it->second);
    }
    xwat.Init(xArray);
    ywat.Init(yArray);
  }
  uint64_t Wat2DSearch::FreqRange2D(uint64_t xs , uint64_t xe , uint64_t ys , uint64_t ye){
    if(xs >= xwat.alphabet_num() || ys >= ywat.alphabet_num())return 0;
    if(xs >= xe || ys >= ye)return 0;
    xe = min(xe , xwat.alphabet_num());
    ye = min(ye , ywat.alphabet_num());
    uint64_t xsPos = xwat.FreqSum(0 , xs);
    uint64_t xePos = xwat.FreqSum(0 , xe);
    if(ye == ywat.alphabet_num()){ // ye >= ywat.alphabet_num()のときNOT_FOUNDが返ってくるため
      return ywat.FreqRange(ys , ye - 1, xsPos , xePos) +
        ywat.Rank(ye - 1 , xePos) - ywat.Rank(ye - 1 , xsPos);
    }
    return ywat.FreqRange(ys , ye , xsPos , xePos);
  }
}

備考

なお、今回の実装は特に論文とか読んだわけではないので、論文とかでやってる二次元探索の方法とは違う可能性があります。たとえばこの実装だと3次元以上のときの拡張とかが難しかったり、wat_arrayを2つ使ってますがもっと少なくてもいいかもしれません。あとx,yの範囲Rが大きければドキュメントにかかれてあることが確かであれば点数によらずO(log R)で計算量が増えるので適宜座標圧縮的なことが必要かもしれません。