DSL_2_C: Range Seach kD Tree

Calendar Clock iconCalendar Clock icon

AOJ

# 目次

# 問題

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

# 解答

# コード

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


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

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

コメント