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

#include <algorithm>
#include <cfloat>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <string>
#include <vector>
typedef long long int ll;
typedef unsigned long long int ull;
#define BIG_NUM 2000000000
#define MOD 1000000007
#define EPS 0.000000001
using namespace std;

struct Node {
  int location;
  int p, left, right;
};

struct Point {
  Point() { id = x = y = 0; }

  Point(int arg_id, int arg_x, int arg_y) {
    id = arg_id;
    x = arg_x;
    y = arg_y;
  }
  bool operator<(const Point &arg) const { return id < arg.id; }
  void print() { printf("%d\\n", id); }
  int id, x, y;
};

static const int MAX = 1000000;
static const int NIL = -1;

int N;
Point P[MAX];
Node T[MAX];
int number;

bool lessX(const Point &p1, const Point &p2) { return p1.x < p2.x; }
bool lessY(const Point &p1, const Point &p2) { return p1.y < p2.y; }

int makeKDTree(int left, int right, int depth) {

  if (!(left < right))
    return NIL;

  int mid = (left + right) / 2;

  int node_id = number++;

  if (depth % 2 == 0) {
    sort(P + left, P + right, lessX);
  } else {
    sort(P + left, P + right, lessY);
  }

  T[node_id].location = mid;
  T[node_id].left = makeKDTree(left, mid, depth + 1);
  T[node_id].right = makeKDTree(mid + 1, right, depth + 1);

  return node_id;
}

void find(int node_id, int sx, int tx, int sy, int ty, int depth,
          vector<Point> &ans) {

  int x = P[T[node_id].location].x;
  int y = P[T[node_id].location].y;

  if (sx <= x && x <= tx && sy <= y && y <= ty) {
    ans.push_back(P[T[node_id].location]);
  }

  if (depth % 2 == 0) {
    if (T[node_id].left != NIL) {
      if (sx <= x)
        find(T[node_id].left, sx, tx, sy, ty, depth + 1, ans);
    }
    if (T[node_id].right != NIL) {
      if (x <= tx)
        find(T[node_id].right, sx, tx, sy, ty, depth + 1, ans);
    }
  } else {
    if (T[node_id].left != NIL) {
      if (sy <= y)
        find(T[node_id].left, sx, tx, sy, ty, depth + 1, ans);
    }
    if (T[node_id].right != NIL) {
      if (y <= ty)
        find(T[node_id].right, sx, tx, sy, ty, depth + 1, ans);
    }
  }
}

int main() {

  int x, y;
  scanf("%d", &N);
  for (int i = 0; i < N; i++) {
    scanf("%d %d", &x, &y);
    P[i] = Point(i, x, y);
    T[i].left = T[i].right = T[i].p = NIL;
  }

  number = 0;

  int root = makeKDTree(0, N, 0);

  int q;
  scanf("%d", &q);
  int sx, tx, sy, ty;
  vector<Point> ans;
  for (int i = 0; i < q; i++) {
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty);
    ans.clear();
    find(root, sx, tx, sy, ty, 0, ans);
    sort(ans.begin(), ans.end());
    for (int j = 0; j < ans.size(); j++) {
      ans[j].print();
    }
    printf("\\n");
  }
  return 0;
}


Error: TLE

Fixed submission:

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

using namespace std;

struct Point {
  int i, x, y;

  Point() {}

  Point(int i, int x, int y) : x(x), y(y), i(i) {}

  void print() { cout << i << endl; }
};
bool lessX(const Point &p1, const Point &p2) { return p1.x < p2.x; }
bool lessY(const Point &p1, const Point &p2) { return p1.y < p2.y; }
bool lessI(const Point &p1, const Point &p2) { return p1.i < p2.i; }

struct Rect {
  int sx, sy, tx, ty;
  Rect(int sx, int sy, int tx, int ty) : sx(sx), sy(sy), tx(tx), ty(ty) {}

  bool cover(Point &p) {
    return (sx <= p.x && tx >= p.x && sy <= p.y && ty >= p.y);
  }
};

class Node {
public:
  Node *ld, *ru;
  Point *p;
};

const int MAX_N = 1000000;
Node T[MAX_N];
int number;

class KdTree {
private:
  Node *root;

public:
  KdTree(Point *pts, int n) {
    sort(pts, pts + n, lessX);
    int l = 0, r = n, m = (l + r) / 2;
    number = 0;
    T[number].p = &pts[m];
    root = T;

    build(root, pts, 0, m, m + 1, r, 1);
  }

  void build(Node *parent, Point *pts, int ll, int lr, int rl, int rr,
             int depth) {
    bool (*c)(const Point &p1, const Point &p2) =
        depth % 2 == 0 ? lessX : lessY;

    // left or down
    if (ll < lr) {
      sort(pts + ll, pts + lr, c);
      int lm = (ll + lr) / 2;
      T[++number].p = &pts[lm];
      parent->ld = &T[number];
      build(parent->ld, pts, ll, lm, lm + 1, lr, depth + 1);
    }

    // right or up
    if (rl < rr) {
      sort(pts + rl, pts + rr, c);
      int rm = (rl + rr) / 2;
      T[++number].p = &pts[rm];
      parent->ru = &T[number];
      build(parent->ru, pts, rl, rm, rm + 1, rr, depth + 1);
    }
  }

  void queryRect(Rect &rect, vector<Point> &inPts) {
    queryRectRecur(rect, root, 0, inPts);
  }

  void queryRectRecur(Rect rect, Node *node, int depth, vector<Point> &inPts) {
    if (node == nullptr) {
      return;
    }

    if (rect.cover(node->p[0]))
      inPts.push_back(node->p[0]);

    if (depth % 2 == 0) {
      if (rect.tx >= node->p[0].x)
        queryRectRecur(rect, node->ru, depth + 1, inPts);
      if (rect.sx <= node->p[0].x)
        queryRectRecur(rect, node->ld, depth + 1, inPts);
    } else {
      if (rect.ty >= node->p[0].y)
        queryRectRecur(rect, node->ru, depth + 1, inPts);
      if (rect.sy <= node->p[0].y)
        queryRectRecur(rect, node->ld, depth + 1, inPts);
    }
  }
};

int main() {
  int n;
  cin >> n;
  Point pts[n];
  for (int i = 0; i < n; ++i) {
    int x, y;
    cin >> x >> y;
    pts[i] = Point(i, x, y);
  }

  KdTree kdTree(pts, n);

  int q;
  cin >> q;
  for (int i = 0; i < q; i++) {
    int sx, tx, sy, ty;
    cin >> sx >> tx >> sy >> ty;
    Rect rect(sx, sy, tx, ty);

    vector<Point> inPts;
    kdTree.queryRect(rect, inPts);

    sort(inPts.begin(), inPts.end(), lessI);
    for (Point p : inPts) {
      p.print();
    }
    cout << endl;
  }

  return 0;
}