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

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pairInt;
typedef tuple<int, int, int> tuple3Int;
#define FOR(i, n) for (int i = 0; i < int(n); i++)
#define FOR1(i, m, n) for (int i = int(m); i < int(n); i++)
#define MAX(a, b) ((a) >= (b) ? (a) : (b))
#define MIN(a, b) ((a) <= (b) ? (a) : (b))

typedef struct _KdNode {
  int id;
  pairInt coords;
  struct _KdNode *parent = nullptr;
  struct _KdNode *left = nullptr;
  struct _KdNode *right = nullptr;
} KdNode;

class KdTree {
private:
  KdNode *root;
  vector<KdNode> nodes;

public:
  KdTree(vector<tuple3Int> &coords) {
    this->nodes.resize(coords.size());
    for (int i = 0; i < coords.size(); i++) {
      nodes[i].id = i;
    }
    this->root = construct(coords, 0, nullptr);
  }

  KdNode *construct(vector<tuple3Int> &coords, int axis, KdNode *parent) {
    if (coords.size() == 0) {
      return nullptr;
    } else if (coords.size() == 1) {
      nodes[get<0>(coords[0])].coords =
          make_pair(get<1>(coords[0]), get<2>(coords[0]));
      nodes[get<0>(coords[0])].parent = parent;
      return &nodes[get<0>(coords[0])];
    }

    vector<tuple3Int>::iterator it;
    vector<vector<tuple3Int>::iterator> temp(10);
    tuple3Int median;
    vector<tuple3Int> left;
    vector<tuple3Int> right;
    const int N = 10;
    if (axis == 0) {
      FOR(i, 10) { temp[i] = coords.begin() + coords.size() * i / 10; }
      sort(temp.begin(), temp.end(),
           [](vector<tuple3Int>::iterator a, vector<tuple3Int>::iterator b) {
             return get<1>(*a) < get<1>(*b);
           });
      it = temp[5];
      median = *it;
      coords.erase(it);
      partition_copy(coords.begin(), coords.end(), back_inserter(left),
                     back_inserter(right),
                     [&](tuple3Int a) { return get<1>(a) < get<1>(median); });
    } else {
      FOR(i, 10) { temp[i] = coords.begin() + coords.size() * i / 10; }
      sort(temp.begin(), temp.end(),
           [](vector<tuple3Int>::iterator a, vector<tuple3Int>::iterator b) {
             return get<2>(*a) < get<2>(*b);
           });
      it = temp[5];
      median = *it;
      coords.erase(it);
      partition_copy(coords.begin(), coords.end(), back_inserter(left),
                     back_inserter(right),
                     [&](tuple3Int a) { return get<2>(a) < get<2>(median); });
    }

    nodes[get<0>(median)].coords = make_pair(get<1>(median), get<2>(median));
    nodes[get<0>(median)].parent = parent;
    nodes[get<0>(median)].left =
        construct(left, (axis + 1) % 2, &nodes[get<0>(median)]);
    nodes[get<0>(median)].right =
        construct(right, (axis + 1) % 2, &nodes[get<0>(median)]);

    return &nodes[get<0>(median)];
  }

  vector<KdNode *> range_query(pairInt a, pairInt b) {
    _v.clear();
    _range_query(root, a, b, 0);
    return _v;
  }

private:
  vector<KdNode *> _v;

  void _range_query(KdNode *node, pairInt a, pairInt b, int axis) {
    if (node == nullptr) {
      return;
    }

    if (axis == 0) {
      if (node->coords.first <= a.first) {
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else if (node->coords.first <= b.first) {
        _range_query(node->left, a, b, (axis + 1) % 2);
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else {
        _range_query(node->left, a, b, (axis + 1) % 2);
      }
    } else {
      if (node->coords.second <= a.second) {
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else if (node->coords.second <= b.second) {
        _range_query(node->left, a, b, (axis + 1) % 2);
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else {
        _range_query(node->left, a, b, (axis + 1) % 2);
      }
    }

    if (a.first <= node->coords.first && a.second <= node->coords.second &&
        node->coords.first <= b.first && node->coords.second <= b.second) {
      _v.push_back(node);
    }

    return;
  }
};

int main(int argc, char *argv[]) {
  int n, q;
  vector<tuple3Int> input;

  scanf("%d", &n);

  FOR(i, n) {
    int x, y;
    scanf("%d%d", &x, &y);
    input.push_back(make_tuple(i, x, y));
  }

  KdTree kd_tree(input);

  scanf("%d", &q);

  FOR(i, q) {
    int x1, x2, y1, y2;
    scanf("%d%d%d%d", &x1, &x2, &y1, &y2);

    vector<KdNode *> v =
        kd_tree.range_query(make_pair(x1, y1), make_pair(x2, y2));
    sort(v.begin(), v.end(),
         [](KdNode *a, KdNode *b) { return a->id < b->id; });

    for (auto it = v.begin(); it != v.end(); it++) {
      printf("%d\\n", (*it)->id);
    }
    printf("\\n");
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> pairInt;
typedef tuple<int, int, int> tuple3Int;
#define FOR(i, n) for (int i = 0; i < int(n); i++)
#define FOR1(i, m, n) for (int i = int(m); i < int(n); i++)
#define MAX(a, b) ((a) >= (b) ? (a) : (b))
#define MIN(a, b) ((a) <= (b) ? (a) : (b))

typedef struct _KdNode {
  int id;
  pairInt coords;
  struct _KdNode *parent = nullptr;
  struct _KdNode *left = nullptr;
  struct _KdNode *right = nullptr;
} KdNode;

class KdTree {
private:
  KdNode *root;
  vector<KdNode> nodes;

public:
  KdTree(vector<tuple3Int> &coords) {
    this->nodes.resize(coords.size());
    for (int i = 0; i < coords.size(); i++) {
      nodes[i].id = i;
    }
    this->root = construct(coords, 0, nullptr);
  }

  KdNode *construct(vector<tuple3Int> &coords, int axis, KdNode *parent) {
    if (coords.size() == 0) {
      return nullptr;
    } else if (coords.size() == 1) {
      nodes[get<0>(coords[0])].coords =
          make_pair(get<1>(coords[0]), get<2>(coords[0]));
      nodes[get<0>(coords[0])].parent = parent;
      return &nodes[get<0>(coords[0])];
    }

    vector<tuple3Int>::iterator it;
    vector<vector<tuple3Int>::iterator> temp(10);
    tuple3Int median;
    vector<tuple3Int> left;
    vector<tuple3Int> right;
    const int N = 10;
    if (axis == 0) {
      FOR(i, 10) { temp[i] = coords.begin() + coords.size() * i / 10; }
      sort(temp.begin(), temp.end(),
           [](vector<tuple3Int>::iterator a, vector<tuple3Int>::iterator b) {
             return get<1>(*a) < get<1>(*b);
           });
      it = temp[5];
      median = *it;
      coords.erase(it);
      partition_copy(coords.begin(), coords.end(), back_inserter(left),
                     back_inserter(right),
                     [&](tuple3Int a) { return get<1>(a) < get<1>(median); });
    } else {
      FOR(i, 10) { temp[i] = coords.begin() + coords.size() * i / 10; }
      sort(temp.begin(), temp.end(),
           [](vector<tuple3Int>::iterator a, vector<tuple3Int>::iterator b) {
             return get<2>(*a) < get<2>(*b);
           });
      it = temp[5];
      median = *it;
      coords.erase(it);
      partition_copy(coords.begin(), coords.end(), back_inserter(left),
                     back_inserter(right),
                     [&](tuple3Int a) { return get<2>(a) < get<2>(median); });
    }

    nodes[get<0>(median)].coords = make_pair(get<1>(median), get<2>(median));
    nodes[get<0>(median)].parent = parent;
    nodes[get<0>(median)].left =
        construct(left, (axis + 1) % 2, &nodes[get<0>(median)]);
    nodes[get<0>(median)].right =
        construct(right, (axis + 1) % 2, &nodes[get<0>(median)]);

    return &nodes[get<0>(median)];
  }

  vector<KdNode *> range_query(pairInt a, pairInt b) {
    _v.clear();
    _range_query(root, a, b, 0);
    return _v;
  }

private:
  vector<KdNode *> _v;

  void _range_query(KdNode *node, pairInt a, pairInt b, int axis) {
    if (node == nullptr) {
      return;
    }

    if (axis == 0) {
      if (node->coords.first <= a.first) {
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else if (node->coords.first <= b.first) {
        _range_query(node->left, a, b, (axis + 1) % 2);
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else {
        _range_query(node->left, a, b, (axis + 1) % 2);
      }
    } else {
      if (node->coords.second <= a.second) {
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else if (node->coords.second <= b.second) {
        _range_query(node->left, a, b, (axis + 1) % 2);
        _range_query(node->right, a, b, (axis + 1) % 2);
      } else {
        _range_query(node->left, a, b, (axis + 1) % 2);
      }
    }

    if (a.first <= node->coords.first && a.second <= node->coords.second &&
        node->coords.first <= b.first && node->coords.second <= b.second) {
      _v.push_back(node);
    }

    return;
  }
};

int main(int argc, char *argv[]) {
  int n, q;
  vector<tuple3Int> input;

  cin >> n;

  FOR(i, n) {
    int x, y;
    cin >> x >> y;
    input.push_back(make_tuple(i, x, y));
  }

  KdTree kd_tree(input);

  cin >> q;

  FOR(i, q) {
    int x1, x2, y1, y2;
    cin >> x1 >> x2 >> y1 >> y2;

    vector<KdNode *> v =
        kd_tree.range_query(make_pair(x1, y1), make_pair(x2, y2));
    sort(v.begin(), v.end(),
         [](KdNode *a, KdNode *b) { return a->id < b->id; });

    for (auto it = v.begin(); it != v.end(); it++) {
      cout << (*it)->id << endl;
    }
    cout << endl;
  }

  return 0;
}