固定のkD-Tree

Calendar Clock iconCalendar Clock icon

tree

# 目次

# Class

// C++ 14
#include <iostream>
#include <string>
#include <vector>
#include <list>
#include <algorithm>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <unordered_map>
#include <math.h>

#define ll long long
#define Int int
#define loop(x, start, end) for(Int x = start; x < end; x++)
#define loopdown(x, start, end) for(int x = start; x > end; x--)
#define rep(n) for(int x = 0; x < n; x++)
#define span(a,x,y) a.begin()+x,a.begin()+y
#define span_all(a) a.begin(),a.end()
#define len(x) (x.size())
#define last(x) (*(x.end()-1))

using namespace std;
#define K 2

typedef struct Point { int coords[K]; int id; } Point;


class KDTree {
public:
    int dimension;
    vector<Point> points;
    
    KDTree(vector<Point> points, int dimension) {
        this->points = points;
        this->dimension = dimension;
        
        buildkDTree();
    }
    
    void searchRange(Point start, Point end, vector<int> &result) {
        searchRangeRec(start, end, result, 0, points.size()-1, 0);
    }
    
private:
    void buildkDTree() {
        buildTreeRec(0, points.size()-1, 0);
    }
    
    bool asc (Point p1, Point p2, int d) {
        return p1.coords[d] < p2.coords[d];
    }

    void buildTreeRec(int left, int right, int depth) {
        if (right < left) return;
        
        int cd = depth % dimension; // current dimension
        auto dAsc = [cd, this](Point p1, Point p2) { return asc(p1, p2, cd); };
        sort(points.begin()+left, points.begin()+right+1, dAsc);
    
        int mid = (left + right) / 2;
        buildTreeRec(left, mid-1, depth+1);
        buildTreeRec(mid+1, right, depth+1);
    }
    
    int checkIncluded(Point point, Point start, Point end, int dimension) {
        if (start.coords[dimension] <= point.coords[dimension] && point.coords[dimension] <= end.coords[dimension]) {
            return 0;
        }
        
        if (point.coords[dimension] < start.coords[dimension]) return -1;
        return 1;
    }
    
    void searchRangeRec(Point start, Point end, vector<int> &result, int left, int right, int depth) {
        if (right < left) return;
        int cd = depth % dimension; // current dimension
        int mid = (left + right) / 2;
        bool inArea = true;
        loop(i,0,dimension) inArea &= checkIncluded(points[mid], start, end, i) == 0;
        if (inArea) result.push_back(points[mid].id);
        if (right == left) return;
        int included = checkIncluded(points[mid], start, end, cd);
        if (included == 0) {
            searchRangeRec(start, end, result, left, mid-1, depth+1);
            searchRangeRec(start, end, result, mid+1, right, depth+1);
        } else if (included == -1) {
            searchRangeRec(start, end, result, mid+1, right, depth+1);
        } else if (included == 1) {
            searchRangeRec(start, end, result, left, mid-1, depth+1);
        }
    }
};

# Usage

// C++ 14
#include <iostream>
#include <string>
#include <vector>
#include <list>
#include <algorithm>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <unordered_map>
#include <math.h>

#define ll long long
#define Int int
#define loop(x, start, end) for(Int x = start; x < end; x++)
#define loopdown(x, start, end) for(int x = start; x > end; x--)
#define rep(n) for(int x = 0; x < n; x++)
#define span(a,x,y) a.begin()+x,a.begin()+y
#define span_all(a) a.begin(),a.end()
#define len(x) (x.size())
#define last(x) (*(x.end()-1))

using namespace std;
int main (void) {
  int n, q;
  scanf("%d", &n);
  vector<Point> points(n);

  loop(i,0,n) {
    cin >> points[i].coords[0] >> points[i].coords[1];
    points[i].id = i;
  }
  KDTree kdTree(points, 2);

  scanf("%d", &q);
  vector<int> result;
  Point start, end;
  loop(i,0,q) {
    result.clear();
    cin >> start.coords[0] >> end.coords[0] >> start.coords[1] >> end.coords[1];
    kdTree.searchRange(start, end, result);
    sort(result.begin(), result.end());

    for (auto id: result) {
      printf("%d\n", id);
    }
    printf("\n");
  }

  return 0;
}

リモートフリーランス。ウェブサービス、スマホアプリエンジニア。
東アジアを拠点に世界を移動しながら活動してます!

お仕事のご依頼・お問い合わせはこちら

コメント