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

#include <algorithm>
#include <cstdio>
#include <vector>

using namespace std;

const int NIL = -1;

struct Node {
  int location = 0;
  int l, r;

  Node(int l = NIL, int r = NIL) : l(l), r(r) {}
};

struct Point {
  int id;
  int x;
  int y;

  Point(int id = 0, int x = 0, int y = 0) : id(id), x(x), y(y) {}

  bool operator<(const Point &other) const { return id < other.id; }
};

bool cmpX(const Point &a, const Point &b) { return a.x < b.x; }

bool cmpY(const Point &a, const Point &b) { return a.y < b.y; }

struct KD {
  int n;
  vector<Point> points;
  vector<Node> heap;
  int np = 0;
  int root = 0;

  KD(const vector<Point> &p) : points(p) {
    n = p.size();
    heap.assign(4 * n, Node());
    root = makeTree(0, n, 0);
  }

  //[l, r)
  int makeTree(int l, int r, int depth) {
    if (l >= r) {
      return NIL;
    }
    int m = l + (r - l) / 2;
    int t = np++;

    if (depth % 2 == 0) {
      sort(points.begin() + l, points.begin() + r, cmpX);
    } else {
      sort(points.begin() + l, points.begin() + r, cmpY);
    }

    heap[t].location = m;
    heap[t].l = makeTree(l, m, depth + 1);
    heap[t].r = makeTree(m + 1, r, depth + 1);

    return t;
  }

  void find(int sx, int tx, int sy, int ty, vector<Point> &ans) {
    find(root, sx, tx, sy, ty, 0, ans);
  }

  void find(int v, int lx, int hx, int ly, int hy, int depth,
            vector<Point> &ans) {

    int x = points[heap[v].location].x;
    int y = points[heap[v].location].y;

    if (lx <= x && x <= hx && ly <= y && y <= hy) {
      ans.push_back(points[heap[v].location]);
      // cout << p[heap[v].locaion].id << " " << p[heap[v].locaion].x << endl;
    }

    if (depth % 2 == 0) {
      if (heap[v].l != NIL && lx <= x) {
        find(heap[v].l, lx, hx, ly, hy, depth + 1, ans);
      }
      if (heap[v].r != NIL && x <= hx) {
        find(heap[v].r, lx, hx, ly, hy, depth + 1, ans);
      }
    }

    else {
      if (heap[v].l != NIL && ly <= y) {
        find(heap[v].l, lx, hx, ly, hy, depth + 1, ans);
      }
      if (heap[v].r != NIL && y <= hy) {
        find(heap[v].r, lx, hx, ly, hy, depth + 1, ans);
      }
    }
  }

  void printStats() {
    for (int i = 0; i < np; i++) {
      Node &node = heap[i];
      //            cout << i << " Loc: " << node.location << " Id: " <<
      //            points[node.location].id
      //                    << " L: " << node.l << " R: " << node.r
      //                    << " (" << points[node.location].x << ","<<
      //                    points[node.location].y << ")"
      //                    << endl;
    }
  }
};

int main() {
  // ios::sync_with_stdio(false);
  int n;
  scanf("%d", &n);
  vector<Point> points;
  for (int i = 0; i < n; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    points.push_back(Point(i, x, y));
  }

  KD kd(points);

  // kd.printStats();

  int q;
  scanf("%d", &q);
  int lx, hx, ly, hy;
  vector<Point> ans;
  for (int i = 0; i < q; i++) {
    scanf("%d%d%d%d", &lx, &hx, &ly, &hy);
    ans.clear();
    kd.find(lx, hx, ly, hy, ans);
    sort(ans.begin(), ans.end());
    for (int j = 0; j < ans.size(); j++) {
      printf("%d\\n", ans[j].id);
    }
    printf("\\n");
  }

  return 0;
}

Error: TLE

Fixed submission:

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

using namespace std;

const int NIL = -1;

struct Point {
  int id;
  int x;
  int y;

  Point(int id = 0, int x = 0, int y = 0) : id(id), x(x), y(y) {}

  bool operator<(const Point &other) const { return id < other.id; }
};

struct Node {
  int location = 0;
  int l, r;
  int minx, maxx, miny, maxy;
  vector<Point> points;

  Node(int l = NIL, int r = NIL) : l(l), r(r) {}
};

bool cmpX(const Point &a, const Point &b) { return a.x < b.x; }

bool cmpY(const Point &a, const Point &b) { return a.y < b.y; }

struct KD {
  int n;
  vector<Point> points;
  vector<Node> heap;
  int np = 0;
  int root = 0;

  KD(const vector<Point> &p) : points(p) {
    n = p.size();
    heap.assign(4 * n, Node());
    sort(points.begin(), points.end(), cmpY);
    root = makeTree(0, n, 0);
  }

  //[l, r)
  int makeTree(int l, int r, int depth) {
    if (l >= r) {
      return NIL;
    }
    int m = l + (r - l) / 2;
    int t = np++;

    // number in points in the rectangle
    heap[t].points.insert(heap[t].points.end(), points.begin() + l,
                          points.begin() + r);

    // heap[t].points = (points.begin() + l, points.begin() + r)
    if (depth % 2 == 0) {
      heap[t].miny = points[l].y;
      heap[t].maxy = points[r - 1].y;
      sort(points.begin() + l, points.begin() + r, cmpX);
      heap[t].minx = points[l].x;
      heap[t].maxx = points[r - 1].x;

    } else {
      heap[t].minx = points[l].x;
      heap[t].maxx = points[r - 1].x;
      sort(points.begin() + l, points.begin() + r, cmpY);
      heap[t].miny = points[l].y;
      heap[t].maxy = points[r - 1].y;
    }

    heap[t].location = m;
    heap[t].l = makeTree(l, m, depth + 1);
    heap[t].r = makeTree(m + 1, r, depth + 1);

    return t;
  }

  void find(int sx, int tx, int sy, int ty, vector<Point> &ans) {
    find(root, sx, tx, sy, ty, 0, ans);
  }

  void find(int v, int qlx, int qhx, int qly, int qhy, int depth,
            vector<Point> &ans) {

    int x = points[heap[v].location].x;
    int y = points[heap[v].location].y;

    if (qlx <= heap[v].minx && heap[v].maxx <= qhx && qly <= heap[v].miny &&
        heap[v].maxy <= qhy) {
      ans.insert(ans.end(), heap[v].points.begin(), heap[v].points.end());
      return;
    }

    // if query cover the range, get the results directly
    if (qlx <= x && x <= qhx && qly <= y && y <= qhy) {
      ans.push_back(points[heap[v].location]);
      // cout << p[heap[v].locaion].id << " " << p[heap[v].locaion].x << endl;
    }

    if (depth % 2 == 0) {
      if (heap[v].l != NIL && qlx <= x) {
        find(heap[v].l, qlx, qhx, qly, qhy, depth + 1, ans);
      }
      if (heap[v].r != NIL && x <= qhx) {
        find(heap[v].r, qlx, qhx, qly, qhy, depth + 1, ans);
      }
    }

    else {
      if (heap[v].l != NIL && qly <= y) {
        find(heap[v].l, qlx, qhx, qly, qhy, depth + 1, ans);
      }
      if (heap[v].r != NIL && y <= qhy) {
        find(heap[v].r, qlx, qhx, qly, qhy, depth + 1, ans);
      }
    }
  }

  void printStats() {
    for (int i = 0; i < np; i++) {
      Node &node = heap[i];
      cout << i << " Loc: " << node.location
           << " Id: " << points[node.location].id << " L: " << node.l
           << " R: " << node.r << " (" << points[node.location].x << ","
           << points[node.location].y << ")" << endl;
    }
  }
};

int main() {
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  vector<Point> points;
  for (int i = 0; i < n; i++) {
    int x, y;
    cin >> x >> y;
    points.push_back(Point(i, x, y));
  }

  KD kd(points);

  // kd.printStats();

  int q;
  cin >> q;
  int lx, hx, ly, hy;
  vector<Point> ans;
  for (int i = 0; i < q; i++) {
    cin >> lx >> hx >> ly >> hy;
    ans.clear();
    kd.find(lx, hx, ly, hy, ans);
    sort(ans.begin(), ans.end());
    for (int j = 0; j < ans.size(); j++) {
      cout << ans[j].id << endl;
    }
    cout << endl;
  }
  return 0;
}