tree

固定のkD-Tree

要素の追加・削除がない固定されたkDTreeを配列で表現する方法を紹介します.
Feb. 2, 2020, 1:51 p.m.

目次

{{ $t('Class') }}

 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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
#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);
        }
    }
};

{{ $t('Usage') }}

 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
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;
}