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

#include <bits/stdc++.h>
using namespace std;

struct Point {
  int x;
  int y;
  int index;
};

struct Node {
  Point p;
  Node *right;
  Node *left;
};

Node *makeKdTree(vector<Point> &ps, int first, int last, int depth) {
  int d = last - first;
  if (d < 0) {
    return nullptr;
  }

  int axis = depth % 2;

  if (axis == 0) {
    sort(ps.begin() + first, ps.begin() + last + 1,
         [](const Point &a, const Point &b) { return a.x < b.x; });

  } else {
    sort(ps.begin() + first, ps.begin() + last + 1,
         [](const Point &a, const Point &b) { return a.y < b.y; });
  }

  Point median = ps[first + d / 2];
  Node *node = new Node;
  node->p = median;
  node->right = makeKdTree(ps, first + d / 2 + 1, last, depth + 1);
  node->left = makeKdTree(ps, first, first + d / 2 - 1, depth + 1);
  return node;
}

void search(Node *p, int sx, int tx, int sy, int ty, int depth,
            vector<int> &result) {
  if (p == nullptr) {
    return;
  }
  if (p->p.x >= sx && p->p.x <= tx && p->p.y >= sy && p->p.y <= ty) {
    result.push_back(p->p.index);
  }

  int axis = depth % 2;

  if (axis == 0) {
    if (p->p.x < sx) {
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    } else if (p->p.x > tx) {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
    } else {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    }
  } else {
    if (p->p.y < sy) {
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    } else if (p->p.y > ty) {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
    } else {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    }
  }

  return;
}

void freeKdTree(Node *p) {
  if (p == nullptr) {
    return;
  }

  freeKdTree(p->left);
  freeKdTree(p->right);
  free(p);
  return;
}

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

  Node *kdtree = makeKdTree(ps, 0, n - 1, 0);

  int q;
  cin >> q;
  for (int i = 0; i < q; i++) {
    int sx, tx, sy, ty;
    cin >> sx >> tx >> sy >> ty;
    vector<int> result;
    search(kdtree, sx, tx, sy, ty, 0, result);
    sort(result.begin(), result.end());
    for (auto item : result) {
      cout << item << '\\n';
    }
    cout << '\\n';
  }

  freeKdTree(kdtree);
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

struct XY {
  int x;
  int y;
  int index;
};

struct Node {
  XY xy;
  Node *right;
  Node *left;
};

void InsertNode(Node *root, XY xy) {
  int depth = 0;
  Node *parent = root;
  Node *child = root;
  bool is_right;
  while (true) {
    if (child == nullptr) {
      Node *node = new Node;
      node->xy = xy;
      node->right = nullptr;
      node->left = nullptr;
      if (is_right) {
        parent->right = node;
      } else {
        parent->left = node;
      }
      break;
    }

    parent = child;
    int axis = depth % 2;
    if (axis == 0) {
      if (child->xy.x < xy.x) {
        child = child->right;
        is_right = true;
      } else {
        child = child->left;
        is_right = false;
      }
    } else {
      if (child->xy.y < xy.y) {
        child = child->right;
        is_right = true;
      } else {
        child = child->left;
        is_right = false;
      }
    }
    depth++;
  }
}

Node *kdTree(vector<XY> v, int first, int last, int depth) {
  int size = last - first;
  if (size < 0) {
    return nullptr;
  }

  int axis = depth % 2;
  if (axis == 0) {
    sort(v.begin() + first, v.begin() + last + 1,
         [](const XY &first, const XY &second) { return first.x < second.x; });
  } else {
    sort(v.begin() + first, v.begin() + last + 1,
         [](const XY &first, const XY &second) { return first.y < second.y; });
  }
  Node *node = new Node;
  node->xy = v[first + size / 2];
  node->right = kdTree(v, first + size / 2 + 1, last, depth + 1);
  node->left = kdTree(v, first, first + size / 2 - 1, depth + 1);

  return node;
}

void search(Node *p, int sx, int tx, int sy, int ty, int depth,
            vector<int> *result) {
  if (p == nullptr) {
    return;
  }

  int axis = depth % 2;

  if (axis == 0) {
    if (p->xy.x < sx) {
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    } else if (p->xy.x > tx) {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
    } else {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
      search(p->right, sx, tx, sy, ty, depth + 1, result);
      if (p->xy.y >= sy && p->xy.y <= ty) {
        result->push_back(p->xy.index);
      }
    }
  } else {
    if (p->xy.y < sy) {
      search(p->right, sx, tx, sy, ty, depth + 1, result);
    } else if (p->xy.y > ty) {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
    } else {
      search(p->left, sx, tx, sy, ty, depth + 1, result);
      search(p->right, sx, tx, sy, ty, depth + 1, result);
      if (p->xy.x >= sx && p->xy.x <= tx) {
        result->push_back(p->xy.index);
      }
    }
  }
}

void freeKdTree(Node *p) {
  if (p == nullptr) {
    return;
  }
  freeKdTree(p->left);
  freeKdTree(p->right);
  free(p);
}

int main() {

  int n;
  cin >> n;
  vector<XY> xys;
  Node *pkdTree = new Node;
  pkdTree->right = nullptr;
  pkdTree->left = nullptr;
  int x, y;
  cin >> x >> y;
  pkdTree->xy = XY{x, y, 0};
  for (int i = 1; i < n; i++) {
    int x, y;
    cin >> x >> y;
    InsertNode(pkdTree, XY{x, y, i});
  }

  // Node *pkdTree = kdTree(xys, 0, n - 1, 0);

  int q;
  cin >> q;
  for (int i = 0; i < q; i++) {
    int sx, tx, sy, ty;
    cin >> sx >> tx >> sy >> ty;
    vector<int> result;
    search(pkdTree, sx, tx, sy, ty, 0, &result);
    sort(result.begin(), result.end());
    for (auto index : result) {
      cout << index << endl;
    }
    cout << endl;
  }

  freeKdTree(pkdTree);
  return 0;
}