概略
C言語でkd treeによるRange search queryを書いた。
コード全体はこちら: https://github.com/fj-th/kd-tree-C-language-
ググってもC言語での実装例のページがあまり出てこないので、一応記事にしておくことに。そこそこわかりやすいコードになっていると(自分では)思います。基本的にこのpdfを参考に実装し、題材とテストとしてAOJ DSL-DSL_2_Cを用いました。木の構築自体はwikipediaの擬似コードもスッキリしていてわかりやすいです。
解きたい問題
Range Query - Range Search (kD Tree)
２次元の平面上の点の集合に対し、与えられた領域に含まれる点を列挙してください。ただし、与えられた点の集合に対して、点の追加・削除は行われません。
要は、x-y平面にばらまかれたデータ点からkd-treeを構成し、その後、指定された長方形領域内に含まれるデータ点を列挙してくださいという問題です。データ自体は一度作ったあと不変でよいので、とにかく範囲探索を早く実現してくれという話ですね。今回は、これをkd-treeを使って解きます。
kd-treeとはなんぞや
wikipediaもありますが、こちらのスライドが視覚的にわかりやすいです。
Range searchまで含めた擬似コードはこのpdfが参考になります(英語ですが)。
kd-treeは、特定軸方向にソートしたデータ点をmedianで分割するということを軸を変えながら繰り返すことで、二分探索に似た発想で高次元データの検索を行えるようにしたデータ構造です。セグメント木の高次元版ですね。
kd-treeは、範囲探索や近傍データ検索が速い一方、データ点の追加や削除、変更を行うと、領域の分割が劇的に変化して木を大規模に作り変える必要が生じる可能性が高いです。そのため、今回の問題のような、静的なデータ集合に繰り返しクエリをかけたいような場合に限り良い選択肢と言えそうです。
実装
木の構築
木を作る部分に関連するコードを抜粋すると以下のような感じで、kdtree()に入力データ点に対応したpoint構造体の配列を与えると、構築したkd-treeのルートノードへのポインタを返すようになっています。
#define XAxis 0
#define YAxis 1
struct point{
int id;
int x;
int y;
};
struct Node{
///// node of kd tree.
struct point location; //// このnodeに対応した分割で使ったpoint
struct Node *leftChild;
struct Node *rightChild;
int depth;
////// leftmost = minimum value of x among the points belonging to the subtree starting from this node.
////// downmost = minimum value of y among the points belonging to the subtree starting from this node.
////// etc...
int leftmost;
int rightmost;
int upmost;
int downmost;
};
struct Node *kdtree(struct point *pointlist, int right, int depth){ ////// [0.....right]についてkd treeを作る。
if (right < 0){
return NULL;
}
if (right == 0){
struct Node *new_node;
new_node = malloc(sizeof(struct Node));
new_node->location = pointlist[right];
new_node->leftChild = NULL;
new_node->rightChild = NULL;
new_node->depth = depth;
new_node->leftmost = new_node->location.x;
new_node->rightmost = new_node->location.x;
new_node->upmost = new_node->location.y;
new_node->downmost = new_node->location.y;
return new_node;
//葉になっているので、子の中の最大値、最小値には自分自身を入れておく。
}
int axis = depth%2; /// 軸が順に選択されるようにする。
struct Node *new_node;
new_node = malloc(sizeof(struct Node));
if (axis == XAxis){
qsort(pointlist,right+1, sizeof(struct point), comparex); ///要素数はright + 1;
new_node->leftmost = pointlist[0].x;
new_node->rightmost = pointlist[right].x; //// x軸方向にカットする場合、qsortの結果を使えばx方向について最大、最小値を取得できる。
}else if (axis == YAxis){
qsort(pointlist,right+1, sizeof(struct point), comparey);
new_node->downmost = pointlist[0].y;
new_node->upmost = pointlist[right].y;
}
int median = (right)/2;
new_node->location = pointlist[median]; ///分割の基準データ点
new_node->depth = depth;
new_node->rightChild = kdtree(pointlist+median+1, right-(median+1), depth +1);
new_node->leftChild = kdtree(pointlist,median, depth +1);
//// y軸方向の最大、最小値をノードと紐づけたい。
//// x軸に切った場合、次の階層ではyで切るので、下のノードでは最大、最小値が確定している。
//// なので、leftChildとrightChildを生成したあとで一つ潜って値を拾ってくる。ただし、NULLの場合のケアに注意。
if (axis == XAxis){
if (new_node->rightChild != NULL && new_node->leftChild != NULL){
new_node->upmost = max(new_node->rightChild->upmost,new_node->leftChild->upmost);
new_node->downmost = min(new_node->rightChild->downmost,new_node->leftChild->downmost);
}else if (new_node->rightChild != NULL){
new_node->upmost = new_node->rightChild->upmost;
new_node->downmost = new_node->rightChild->downmost;
}else if (new_node->leftChild != NULL){
new_node->upmost = new_node->leftChild->upmost;
new_node->downmost = new_node->leftChild->downmost;
}else{////葉だった場合。
new_node->upmost = new_node->location.y;
new_node->downmost = new_node->location.y;
}
}else{
if (new_node->rightChild != NULL && new_node->leftChild != NULL){
new_node->rightmost = max(new_node->rightChild->rightmost,new_node->leftChild->rightmost);
new_node->leftmost = min(new_node->rightChild->leftmost,new_node->leftChild->leftmost);
}else if (new_node->rightChild != NULL){
new_node->rightmost = new_node->rightChild->rightmost;
new_node->leftmost = new_node->rightChild->leftmost;
}else if (new_node->leftChild != NULL){
new_node->rightmost = new_node->leftChild->rightmost;
new_node->leftmost = new_node->leftChild->leftmost;
}else{
new_node->rightmost = new_node->location.x;
new_node->leftmost = new_node->location.x;
}
}
return new_node;
}
ここで、kd treeのノードには、そのノードが表す分割点を表すlocationや、そのノードのkd tree内での深さを表すdepth、そして、そのノードを根とする部分木に属する点のx, y座標の最大・最小値を入れています。
x軸とy軸交互に見て、残っているデータ点のmedianで分割を行い、新しいノードを作成します。各ノードに、それをルートとする部分木に含まれるデータのx, y座標の最大・最小値をもたせており、これにより、以下のようにRange searchを行うときの探索・枝刈りが高速に行えます。
Range search
int is_contained(struct Node *region, int sx, int tx, int sy, int ty){
///// regionノードから始まるsubtreeの要素が指定領域にすっぽり収まっているか否かを返す。
///// 各ノードに、自分以下の子の最大、最小値をもたせているので、それを参照するだけで判定できる。
int leftmost = region->leftmost;
int rightmost = region->rightmost;
int upmost = region->upmost;
int downmost = region->downmost;
if (leftmost < sx || rightmost > tx || upmost > ty || downmost < sy) return 0;
return 1;
}
void SearchKDtree(struct Node *v, int sx, int tx, int sy, int ty){
///// 探索の枝刈りをどう行うか。
///// ノードvを根とするsubtree中のxの最大値がsxより小さいならば、v以下の探索は行わなくても良い, etc...という感じ。
if (v->rightmost < sx || v->leftmost > tx || v->downmost > ty || v->upmost < sy) return; ///枝刈り部分
if (v->leftChild == NULL && v->rightChild == NULL){
if (sx <= v->location.x && sy <= v->location.y && tx >= v->location.x && ty >= v->location.y){
id_collection[leftcur] = (v->location).id;
leftcur +=1;
}
return; //vが葉なら領域内か判定
}
if (v->leftChild != NULL){ //子がいれば潜る
if (is_contained(v->leftChild, sx, tx, sy, ty)){
ReportSubtree(v->leftChild);
}else{
SearchKDtree(v->leftChild,sx, tx, sy, ty);
}
}
if (v->rightChild != NULL){
if (is_contained(v->rightChild, sx, tx, sy, ty)){
ReportSubtree(v->rightChild);
}else{
SearchKDtree(v->rightChild,sx, tx, sy, ty);
}
}
return;
}
ここで、ReportSubtree()は、subtreeが表す領域内のデータ点を全て列挙するような関数で、is_contained()がTrueなときに走ります。SearchKDtree()は、グローバルに置いたid_collectionに、範囲探索で該当したデータ点を詰めています。id_collectionの要素数はleftcurで管理しています。最終的には、id_collectionの[0...leftcur)が答えとなります(idについて昇順にするならばあとでソートする)。
ノードに部分木のx, y座標の最大・最小値をもたせるなど、データ構造自体をリッチにすることで、処理を早くしています。逆に言えば、先に述べたとおり、データ自体に変更が加わってしまうとその影響は一般に甚大です。