aoj

DSL_2_C: Range Seach kD Tree

DSL_2_C: Range Seach kD Tree
Feb. 2, 2020, 1:52 p.m.

目次

問題

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_C

解答

コード

 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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
#define K 2


typedef struct Node {
  int id;
  int point[K];
  Node *left;
  Node *right;
} Node;

Node *newNode(int id, int point[K]) {
  Node *node = new Node;
  node->id = id;
  loop(i,0,K) node->point[i] = point[i];
  node->left = node->right = NULL;
  return node;
}


Node* insertRec(Node *root, int id, int point[K], int depth) {
  if (root == NULL) return newNode(id, point);

  int cd = depth % K;
  if (point[cd] < root->point[cd])
    root->left = insertRec(root->left, id, point, depth+1);
  if (root->point[cd] <= point[cd])
    root->right = insertRec(root->right, id, point, depth+1);
  return root;
}


Node *insert(Node *root, int id, int point[K]) {
  return insertRec(root, id, point, 0);
}


void rangeSearch(Node *root, int start[K], int end[K], int depth, vector<int> &result) {
  if (root == NULL) return;

  bool inArea = true;
  loop(i,0,K) {
    inArea &= start[i] <= root->point[i];
    inArea &= root->point[i] <= end[i];
  }

  if (inArea) {
    result.push_back(root->id);
  }

  int cd = depth % K;
  int p = root->point[cd], s = start[cd], e = end[cd];

  // s   p    e
  // ----------
  if (s <= p && p <= e) {
    rangeSearch(root->right, start, end, depth+1, result);
    rangeSearch(root->left, start, end, depth+1, result);
    return;
  }

  //    p  s       e            s      e  p
  // i) ---------------   ii) --------------
  if (p < s) {
    rangeSearch(root->right, start, end, depth+1, result);
  } else {
    rangeSearch(root->left, start, end, depth+1, result);
  }
}


int main(void) {
  int n, q;
  cin >> n;
  Node *root = NULL;
  loop(i,0,n) {
    int point[K];
    cin >> point[0] >> point[1];
    root = insert(root, i, point);
  }

  cin >> q;
  loop(i,0,q) {
    int start[K], end[K];
    cin >> start[0] >> end[0] >> start[1] >> end[1];
    vector<int> result;
    rangeSearch(root, start, end, 0, result);
    for (auto id: result) {
      cout << id << endl;
    }
    cout << endl;
  }

  return 0;
}