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

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

#define rep(i, x, y) for (int i = (x); i < (y); ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define print(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define print(x)
#endif

const int inf = 1e9;
const int64_t inf64 = 1e18;
const double eps = 1e-9;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

using i64 = int64_t;

template <int D, class point> struct kd_tree {
  point p;
  int axis;
  kd_tree<D, point> *left, *right;
  kd_tree(int axis_ = 0) : left(NULL), right(NULL), axis(axis_){};
  kd_tree(vector<point> points, int axis_ = 0)
      : left(NULL), right(NULL), axis(axis_) {
    build(points, 0, points.size());
  }
  ~kd_tree() {
    delete left;
    delete right;
  }
  kd_tree<D, point> *build(vector<point> &points, int l, int r) {
    int m = (l + r) / 2;
    nth_element(
        points.begin() + l, points.begin() + m, points.begin() + r,
        [&](const point &p, const point &q) { return p[axis] < q[axis]; });
    p = points[m];
    if (l < m) {
      left = new kd_tree<D, point>((axis + 1) % D);
      left->build(points, l, m);
    }
    if (m + 1 < r) {
      right = new kd_tree<D, point>((axis + 1) % D);
      right->build(points, m + 1, r);
    }
    return this;
  }
  kd_tree<D, point> *insert(const point &q) {
    if (q[axis] < p[axis]) {
      if (left == NULL) {
        vector<point> points = {q};
        left = new kd_tree<D, point>(points, (axis + 1) % D);
        return left;
      } else
        return left->insert(q);
    } else {
      if (right == NULL) {
        vector<point> points = {q};
        left = new kd_tree<D, point>(points, (axis + 1) % D);
        return left;
      } else
        return right->insert(q);
    }
  }
  //????????????(O(n))?????§erased?????°?????????????????????,erased???????????\\????¢????????????¨??????build???????????????????????????????????????.
  void erase(const point &q) {
    kd_tree<D, point> *tree = nearest_neighbor(q);
    vector<point> points = tree->collect();
    points.erase(points.rbegin());
    delete tree->left;
    delete tree->right;
    tree->build(points, 0, points.size());
  }
  vector<point> collect() const {
    vector<point> points;
    if (left != NULL) {
      vector<point> left_points = left->collect();
      points.insert(points.end(), left_points.begin(), left_points.end());
    }
    if (right != NULL) {
      vector<point> right_points = right->collect();
      points.insert(points.end(), right_points.begin(), right_points.end());
    }
    points.push_back(p);
    return points;
  };
  kd_tree<D, point> *nearest_neighbor(const point &q) const {
    kd_tree<D, point> *res = this;
    long double mind = inf64;
    auto dist = [&](const point &p, const point &q) {
      long double d = 0;
      for (int i = 0; i < D; ++i)
        d += (p[i] - q[i]) * (p[i] - q[i]);
      return sqrtl(d);
    };
    function<void(const kd_tree<D, point> *)> rec =
        [&](const kd_tree<D, point> *tree) {
          if (dist(p, q) < mind) {
            res = tree;
            mind = dist(p, q);
          }
          if (q[tree->axis] < tree->p[tree->axis]) {
            if (tree->left != NULL)
              rec(tree->left);
            if (tree->p[tree->axis] - q[tree->axis] <= mind)
              rec(tree->right);
          } else {
            if (tree->right != NULL)
              rec(tree->right);
            if (q[tree->axis] - tree->p[tree->axis] <= mind)
              rec(tree->left);
          }
        };
    rec(this);
    return res;
  }
  vector<point> range_search(const point &lb, const point &ub) const {
    vector<point> res;
    function<void(const kd_tree<D, point> *)> rec =
        [&](const kd_tree<D, point> *tree) {
          bool in_range = true;
          for (int i = 0; i < D; ++i)
            if (tree->p[i] < lb[i] or ub[i] < tree->p[i]) {
              in_range = false;
              break;
            }
          if (in_range)
            res.push_back(tree->p);
          if (tree->left != NULL and lb[tree->axis] <= tree->p[tree->axis])
            rec(tree->left);
          if (tree->right != NULL and ub[tree->axis] >= tree->p[tree->axis])
            rec(tree->right);
        };
    rec(this);
    return res;
  }
};

void solve() {
  int n;
  scanf("%d", &n);
  vector<vector<int>> ps(n);
  unordered_map<int, unordered_map<int, vector<int>>> mp;
  rep(i, 0, n) {
    int x, y;
    scanf("%d %d", &x, &y);
    ps[i] = {x, y};
    mp[x][y].push_back(i);
  }
  kd_tree<2, vector<int>> kt(ps);
  int q;
  scanf("%d", &q);
  rep(i, 0, q) {
    int sx, tx, sy, ty;
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty);
    vector<int> lb = {sx, sy}, ub = {tx, ty};
    auto tmp = kt.range_search(lb, ub);
    vector<int> ans;
    for (auto p : tmp)
      ans.insert(ans.end(), mp[p[0]][p[1]].begin(), mp[p[0]][p[1]].end());
    sort(ans.begin(), ans.end());
    rep(j, 0, ans.size()) printf("%d\\n", ans[j]);
    printf("\\n");
  }
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(10);
  solve();
  return 0;
}

Error: TLE

Fixed submission:

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

#define rep(i, x, y) for (int i = (x); i < (y); ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define print(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define print(x)
#endif

const int inf = 1e9;
const int64_t inf64 = 1e18;
const double eps = 1e-9;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

using i64 = int64_t;

template <int D, class point> struct kd_tree {
  point p;
  int axis;
  kd_tree<D, point> *left, *right;
  kd_tree(int axis_) : axis(axis_){};
  kd_tree(vector<point> points, int axis_ = 0)
      : left(NULL), right(NULL), axis(axis_) {
    build(points, 0, points.size());
  }
  ~kd_tree() {
    delete left;
    delete right;
  }
  kd_tree<D, point> *build(vector<point> &points, int l, int r) {
    int m = (l + r) / 2;
    sort(points.begin() + l, points.begin() + r,
         [&](const point &p, const point &q) { return p[axis] < q[axis]; });
    p = points[m];
    if (l < m) {
      left = new kd_tree<D, point>((axis + 1) % D);
      left->build(points, l, m);
    }
    if (m + 1 < r) {
      right = new kd_tree<D, point>((axis + 1) % D);
      right->build(points, m + 1, r);
    }
    return this;
  }
  kd_tree<D, point> *insert(const point &q) {
    if (q[axis] < p[axis]) {
      if (left == NULL) {
        vector<point> points = {q};
        left = new kd_tree<D, point>(points, (axis + 1) % D);
        return left;
      } else
        return left->insert(q);
    } else {
      if (right == NULL) {
        vector<point> points = {q};
        left = new kd_tree<D, point>(points, (axis + 1) % D);
        return left;
      } else
        return right->insert(q);
    }
  }
  //????????????(O(n))?????§erased?????°?????????????????????,erased???????????\\????¢????????????¨??????build???????????????????????????????????????.
  void erase(const point &q) {
    kd_tree<D, point> *tree = nearest_neighbor(q);
    vector<point> points = tree->collect();
    points.erase(points.rbegin());
    delete tree->left;
    delete tree->right;
    tree->build(points);
  }
  vector<point> collect() const {
    vector<point> points;
    if (left != NULL) {
      vector<point> left_points = left->collect();
      points.insert(points.end(), left_points.begin(), left_points.end());
    }
    if (right != NULL) {
      vector<point> right_points = right->collect();
      points.insert(points.end(), right_points.begin(), right_points.end());
    }
    points.push_back(p);
    return points;
  };
  kd_tree<D, point> *nearest_neighbor(const point &q) const {
    kd_tree<D, point> *res = this;
    long double mind = inf64;
    auto dist = [&](const point &p, const point &q) {
      long double d = 0;
      for (int i = 0; i < D; ++i)
        d += (p[i] - q[i]) * (p[i] - q[i]);
      return sqrtl(d);
    };
    function<void(const kd_tree<D, point> *)> rec =
        [&](const kd_tree<D, point> *tree) {
          if (dist(p, q) < mind) {
            res = tree;
            mind = dist(p, q);
          }
          if (q[tree->axis] < tree->p[tree->axis]) {
            if (tree->left != NULL)
              rec(tree->left);
            if (tree->p[tree->axis] - q[tree->axis] <= mind)
              rec(tree->right);
          } else {
            if (tree->right != NULL)
              rec(tree->right);
            if (q[tree->axis] - tree->p[tree->axis] <= mind)
              rec(tree->left);
          }
        };
    rec(this);
    return res;
  }
  vector<point> range_search(const point &lb, const point &ub) const {
    vector<point> res;
    function<void(const kd_tree<D, point> *)> rec =
        [&](const kd_tree<D, point> *tree) {
          bool in_range = true;
          for (int i = 0; i < D; ++i)
            if (tree->p[i] < lb[i] or ub[i] < tree->p[i]) {
              in_range = false;
              break;
            }
          if (in_range)
            res.push_back(tree->p);
          if (tree->left != NULL and lb[tree->axis] <= tree->p[tree->axis])
            rec(tree->left);
          if (tree->right != NULL and ub[tree->axis] >= tree->p[tree->axis])
            rec(tree->right);
        };
    rec(this);
    return res;
  }
};

void solve() {
  int n;
  scanf("%d", &n);
  vector<vector<int>> ps(n);
  unordered_map<int, unordered_map<int, vector<int>>> mp;
  rep(i, 0, n) {
    int x, y;
    scanf("%d %d", &x, &y);
    ps[i] = {x, y};
    mp[x][y].push_back(i);
  }
  kd_tree<2, vector<int>> kt(ps);
  int q;
  scanf("%d", &q);
  rep(i, 0, q) {
    int sx, tx, sy, ty;
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty);
    vector<int> lb = {sx, sy}, ub = {tx, ty};
    auto tmp = kt.range_search(lb, ub);
    vector<int> ans;
    for (auto p : tmp)
      ans.insert(ans.end(), mp[p[0]][p[1]].begin(), mp[p[0]][p[1]].end());
    sort(ans.begin(), ans.end());
    rep(j, 0, ans.size()) printf("%d\\n", ans[j]);
    printf("\\n");
  }
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(10);
  solve();
  return 0;
}