Problem p02347 (C++) — buggy submission, verdict: Time Limit Exceeded

//
// Created by tattsun on 2017/02/21.
//

#include <algorithm>
#include <iostream>
#include <vector>

#define POINT_MAX 500001

using namespace std;

int n, q;
int xs[POINT_MAX], ys[POINT_MAX];

class Point {
public:
  Point(int id, int x, int y) : id(id), x(x), y(y) {}
  int id, x, y;
};

class Node {
public:
  Node() : l(), r(), location(-1) {}
  Node *l;
  Node *r;
  int location;
};

bool ascX(const Point &x, const Point &y) { return x.x < y.x; }

bool ascY(const Point &x, const Point &y) { return x.y < y.y; }

class KdTree {
public:
  KdTree(vector<Point> &ps) : ps(ps) {
    root = buildKdTree(ps, 0, (int)ps.size(), 0);
  }

  void find(int x1, int x2, int y1, int y2, vector<int> &res) {
    _find(x1, x2, y1, y2, 0, root, res);
    sort(res.begin(), res.end());
  }

private:
  vector<Point> &ps;
  Node *root;

  void _find(int x1, int x2, int y1, int y2, int depth, Node *n,
             vector<int> &res) {
    Point &p(ps[n->location]);
    if (x1 <= p.x && p.x <= x2 && y1 <= p.y && p.y <= y2) {
      res.push_back(p.id);
    }

    if (depth % 2 == 0) {
      if (n->l != nullptr && p.x >= x1) {
        _find(x1, x2, y1, y2, depth + 1, n->l, res);
      }
      if (n->r != nullptr && p.x <= x2) {
        _find(x1, x2, y1, y2, depth + 1, n->r, res);
      }
    } else {
      if (n->l != nullptr && p.y >= y1) {
        _find(x1, x2, y1, y2, depth + 1, n->l, res);
      }
      if (n->r != nullptr && p.y <= y2) {
        _find(x1, x2, y1, y2, depth + 1, n->r, res);
      }
    }
  }

  static Node *buildKdTree(vector<Point> &ps, int l, int r, int depth) {
    if (l >= r) {
      return nullptr;
    }

    if (depth % 2 == 0) {
      sort(ps.begin() + l, ps.begin() + r, ascX);
    } else {
      sort(ps.begin() + l, ps.begin() + r, ascY);
    }

    int mid = (l + r) / 2;

    Node *n = new Node();
    n->location = mid;
    n->l = buildKdTree(ps, l, mid, depth + 1);
    n->r = buildKdTree(ps, mid + 1, r, depth + 1);

    return n;
  }
};

int main() {
  vector<Point> ps;

  scanf("%d", &n);

  int x, y;
  for (int i = 0; i < n; i++) {
    scanf("%d %d", &x, &y);
    ps.push_back(Point(i, x, y));
  }

  KdTree kd(ps);

  scanf("%d", &q);

  int x1, x2, y1, y2;
  for (int i = 0; i < q; i++) {
    vector<int> ans;
    scanf("%d %d %d %d", &x1, &x2, &y1, &y2);
    kd.find(x1, x2, y1, y2, ans);
    for (auto it = ans.begin(); it != ans.end(); it++) {
      printf("%d\\n", *it);
    }
    printf("\\n");
  }
}

Error: TLE

Fixed submission:

//
// Created by tattsun on 2017/02/21.
//

#include <algorithm>
#include <iostream>
#include <vector>

#define POINT_MAX 500001

using namespace std;

int n, q;
int xs[POINT_MAX], ys[POINT_MAX];

class Point {
public:
  Point(int id, int x, int y) : id(id), x(x), y(y) {}
  int id, x, y;
};

class Node {
public:
  Node() : l(), r(), location(-1) {}
  Node *l;
  Node *r;
  int location;
};

bool ascX(const Point &x, const Point &y) { return x.x < y.x; }

bool ascY(const Point &x, const Point &y) { return x.y < y.y; }

class KdTree {
public:
  KdTree(vector<Point> &ps) : ps(ps) {
    root = buildKdTree(ps, 0, (int)ps.size(), 0);
  }

  void find(int x1, int x2, int y1, int y2, vector<int> &res) {
    _find(x1, x2, y1, y2, 0, root, res);
    sort(res.begin(), res.end());
  }

private:
  vector<Point> &ps;
  Node *root;

  void _find(int x1, int x2, int y1, int y2, int depth, Node *n,
             vector<int> &res) {
    Point &p(ps[n->location]);
    if (x1 <= p.x && p.x <= x2 && y1 <= p.y && p.y <= y2) {
      res.push_back(p.id);
    }

    if (depth % 2 == 0) {
      if (n->l != nullptr && p.x >= x1) {
        _find(x1, x2, y1, y2, depth + 1, n->l, res);
      }
      if (n->r != nullptr && p.x <= x2) {
        _find(x1, x2, y1, y2, depth + 1, n->r, res);
      }
    } else {
      if (n->l != nullptr && p.y >= y1) {
        _find(x1, x2, y1, y2, depth + 1, n->l, res);
      }
      if (n->r != nullptr && p.y <= y2) {
        _find(x1, x2, y1, y2, depth + 1, n->r, res);
      }
    }
  }

  static Node *buildKdTree(vector<Point> &ps, int l, int r, int depth) {
    if (l >= r) {
      return nullptr;
    }

    if (depth % 2 == 0) {
      sort(ps.begin() + l, ps.begin() + r, ascX);
    } else {
      sort(ps.begin() + l, ps.begin() + r, ascY);
    }

    int mid = (l + r) / 2;

    Node *n = new Node();
    n->location = mid;
    n->l = buildKdTree(ps, l, mid, depth + 1);
    n->r = buildKdTree(ps, mid + 1, r, depth + 1);

    return n;
  }
};

int main() {
  vector<Point> ps;

  cin >> n;

  int x, y;
  for (int i = 0; i < n; i++) {
    cin >> x >> y;
    ps.push_back(Point(i, x, y));
  }

  KdTree kd(ps);

  cin >> q;

  int x1, x2, y1, y2;
  for (int i = 0; i < q; i++) {
    vector<int> ans;
    cin >> x1 >> x2 >> y1 >> y2;
    kd.find(x1, x2, y1, y2, ans);
    for (auto it = ans.begin(); it != ans.end(); it++) {
      cout << *it << endl;
    }
    cout << endl;
  }
}