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

#include <iostream>
using namespace std;
#include <algorithm>
#include <cstdio> // k-D_tree. (k = 2)

class Point {
public:
  int id, x, y;
  Point(){};
  Point(int no, int pos_x, int pos_y) : id(no), x(pos_x), y(pos_y){};
  void print() { printf("%d\\n", id); }
};

bool cmp_x(const Point &P1, const Point &P2) { return P1.x < P2.x; }
bool cmp_y(const Point &P1, const Point &P2) { return P1.y < P2.y; }
bool cmp_id(const Point &P1, const Point &P2) { return P1.id < P2.id; }

class node // nodes of 2-D_tree.
{
public:
  Point *mark;
  node *less, *more;
  int depth; // parity of depth is important.
  node() : mark(NULL), less(NULL), more(NULL), depth(0){};
  node(Point *pos, node *l, node *m, int d)
      : mark(pos), less(l), more(m), depth(d){};
};

node *make_tree(Point S[], int left, int right,
                int d) // recursive construction of 2-D_tree.
{
  if (left > right) {
    return NULL;
  }
  if (left == right) {
    node *leaf = new node;
    *leaf = node(&S[left], NULL, NULL, d); // arrived at leaf.
    return leaf;
  }

  if (d & 1) {
    sort(S + left, S + right + 1, cmp_x);
  } else {
    sort(S + left, S + right + 1, cmp_y);
  }

  int mid = (left + right) / 2;

  node *n = new node; // make a new node.
  n->mark = &S[mid];
  n->depth = d;
  n->less = make_tree(S, left, mid - 1, d + 1);
  n->more = make_tree(S, mid + 1, right, d + 1);

  return n;
}

int cnt;

void find(Point S[], int sx, int tx, int sy, int ty, node *kd_n) {
  if (kd_n == NULL)
    return;

  if (sx <= kd_n->mark->x && kd_n->mark->x <= tx && sy <= kd_n->mark->y &&
      kd_n->mark->y <= ty) {
    S[cnt] = *(kd_n->mark);
    cnt++;
  }
  int ax = kd_n->mark->x, ay = kd_n->mark->y, d = kd_n->depth;

  if (d & 1) {
    if (sx <= ax)
      find(S, sx, tx, sy, ty, kd_n->less);
    if (ax <= tx)
      find(S, sx, tx, sy, ty, kd_n->more);
  } else {
    if (sy <= ay)
      find(S, sx, tx, sy, ty, kd_n->less);
    if (ay <= ty)
      find(S, sx, tx, sy, ty, kd_n->more);
  }
}

int main() {
  int i, pos_x, pos_y, n;
  Point P[500000], Q[100];

  scanf("%d", &n);
  for (i = 0; i < n; i++) {
    scanf("%d %d", &pos_x, &pos_y);
    P[i] = Point(i, pos_x, pos_y);
  }

  node *root;
  root = make_tree(P, 0, n - 1, 0); // make a 2-D_tree.

  int q, sx, tx, sy, ty;
  scanf("%d", &q);
  while (q) {
    cnt = 0;
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty); // input data.
    find(Q, sx, tx, sy, ty, root);            // counting.
    sort(Q, Q + cnt, cmp_id);
    for (i = 0; i < cnt; i++) {
      Q[i].print();
    } // print out.
    cout << endl;
    q--;
  };

  return 0;
}

Error: TLE

Fixed submission:

#include <iostream>
using namespace std;
#include <stdio.h> // k-D_tree. (k = 2)

class Point {
public:
  int id, x, y;
  Point(){};
  Point(int no, int pos_x, int pos_y) : id(no), x(pos_x), y(pos_y){};
  void print() { printf("%d\\n", id); }
};

void swap_p(Point &P1, Point &P2) {
  swap(P1.x, P2.x);
  swap(P1.y, P2.y);
  swap(P1.id, P2.id);
}

int part_ax(Point S[], int left, int right, bool axis, bool idswap = false) {
  int i = left - 1, j;
  int val = (idswap ? S[right].id : (axis ? S[right].x : S[right].y));

  for (j = left; j < right; j++) {
    if ((idswap ? S[j].id : (axis ? S[j].x : S[j].y)) <= val) {
      i++;
      swap_p(S[i], S[j]);
    }
  }
  swap_p(S[i + 1], S[right]);

  return i + 1;
}

void qsort_ax(Point S[], int left, int right, bool axis, bool idswap = false) {
  if (!(left < right))
    return;
  int h;
  h = part_ax(S, left, right, axis, idswap);
  qsort_ax(S, left, h - 1, axis, idswap);
  qsort_ax(S, h + 1, right, axis, idswap);
}

class node // nodes of 2-D_tree.
{
public:
  Point *mark;
  node *less, *more;
  int depth; // parity of depth is important.
  node() : mark(NULL), less(NULL), more(NULL), depth(0){};
  node(Point *pos, node *l, node *m, int d)
      : mark(pos), less(l), more(m), depth(d){};
};

node *make_tree(Point S[], int left, int right,
                int d) // recursive construction of 2-D_tree.
{
  if (left > right) {
    return NULL;
  }
  if (left == right) {
    node *leaf = new node;
    *leaf = node(&S[left], NULL, NULL, d); // arrived at leaf.
    return leaf;
  }
  qsort_ax(S, left, right, d & 1); // pick up the median.

  int mid = (left + right) / 2;

  node *n = new node; // make a new node.
  n->mark = &S[mid];
  n->depth = d;
  n->less = make_tree(S, left, mid - 1, d + 1);
  n->more = make_tree(S, mid + 1, right, d + 1);

  return n;
}

int count;

void find(Point S[], int sx, int tx, int sy, int ty, node *kd_n) {
  if (kd_n == NULL)
    return;

  if (sx <= kd_n->mark->x && kd_n->mark->x <= tx && sy <= kd_n->mark->y &&
      kd_n->mark->y <= ty) {
    S[count] = *(kd_n->mark);
    count++;
  }
  int ax = kd_n->mark->x, ay = kd_n->mark->y, d = kd_n->depth;

  if (d & 1) {
    if (sx <= ax)
      find(S, sx, tx, sy, ty, kd_n->less);
    if (ax <= tx)
      find(S, sx, tx, sy, ty, kd_n->more);
  } else {
    if (sy <= ay)
      find(S, sx, tx, sy, ty, kd_n->less);
    if (ay <= ty)
      find(S, sx, tx, sy, ty, kd_n->more);
  }
}

int main() {
  int i, pos_x, pos_y, n;
  Point P[500000], Q[100];

  scanf("%d", &n);
  for (i = 0; i < n; i++) {
    scanf("%d %d", &pos_x, &pos_y);
    P[i] = Point(i, pos_x, pos_y);
  }

  node *root;
  root = make_tree(P, 0, n - 1, 0); // make a 2-D_tree.

  int q, sx, tx, sy, ty;
  scanf("%d", &q);
  while (q) {
    count = 0;
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty); // input data.
    find(Q, sx, tx, sy, ty, root);            // counting.
    qsort_ax(Q, 0, count - 1, 0, 1);          // sorting.
    for (i = 0; i < count; i++) {
      Q[i].print();
    } // print out.
    cout << endl;
    q--;
  };

  return 0;
}