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

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

using namespace std;

struct Point {
  int id;
  int x;
  int y;

  Point(int id, int x, int y) : id(id), x(x), y(y) {}
  Point() : Point(0, 0, 0) {}
};

bool lessID(const Point &p1, const Point &p2) { return p1.id < p2.id; }
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; }

using PointList = vector<Point>;

class KdTree {
  struct Node {
    // int depth;
    Point location;
    unique_ptr<Node> left;
    unique_ptr<Node> right;
  };

  using NodePtr = unique_ptr<Node>;

  NodePtr build(PointList::iterator first, PointList::iterator last,
                int depth) {
    auto k = distance(first, last);
    if (k < 1)
      return nullptr;

    auto node = make_unique<Node>();
    if (k == 1) {
      node->location = *first;
      return node;
    }

    if (depth % 2 == 0) {
      // horizontal
      sort(first, last, lessX);
    } else {
      // vertical
      sort(first, last, lessY);
    }
    auto mid = first + distance(first, last) / 2;
    node->location = *mid;
    node->left = build(first, mid, depth + 1);
    node->right = build(mid + 1, last, depth + 1);
    return node;
  }

public:
  explicit KdTree(PointList &points) : _root(nullptr) {
    _root = build(points.begin(), points.end(), 0);
  }

  PointList find(int sx, int sy, int tx, int ty) {
    PointList points;
    find(*_root, sx, sy, tx, ty, 0, points);
    sort(points.begin(), points.end(), lessID);
    return points;
  }

  void find(const Node &node, int sx, int sy, int tx, int ty, int depth,
            PointList &ans) {
    auto p = node.location;
    if (sx <= p.x && p.x <= tx && sy <= p.y && p.y <= ty) {
      // cerr << "query: " << sx << " " << sy << " " << tx << " " << ty << endl;
      // cerr << "point(" << p.id << "): " << p.x << ", " << p.y << endl;
      // cerr << "debug: " << (sx <= p.x) << " " << (p.x <= tx) << " ";
      // cerr << (sy <= p.y) << " " << (p.y <= ty) << endl;
      ans.push_back(p);
    }

    if (depth % 2 == 0) {
      // horizontal
      if (node.left != nullptr && sx <= p.x) {
        find(*node.left, sx, sy, tx, ty, depth + 1, ans);
      }
      if (node.right != nullptr && p.x <= tx) {
        find(*node.right, sx, sy, tx, ty, depth + 1, ans);
      }
    } else {
      // vertical
      if (node.left != nullptr && sy <= p.y) {
        find(*node.left, sx, sy, tx, ty, depth + 1, ans);
      }
      if (node.right != nullptr && p.y <= ty) {
        find(*node.right, sx, sy, tx, ty, depth + 1, ans);
      }
    }
  }

private:
  NodePtr _root;
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  PointList points;
  for (int i = 0; i < n; ++i) {
    int x, y;
    cin >> x >> y;
    points.emplace_back(i, x, y);
  }
  KdTree kdtree(points);

  int q;
  cin >> q;
  for (int i = 0; i < q; ++i) {
    int sx, sy, tx, ty;
    cin >> sx >> tx >> sy >> ty;
    auto found = kdtree.find(sx, sy, tx, ty);
    for (auto &p : found) {
      cout << p.id << "\\n";
    }
    cout << "\\n";
  }

  return 0;
}

Error: TLE

Fixed submission:

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

using namespace std;

struct Point {
  int id;
  int x;
  int y;

  Point(int id, int x, int y) : id(id), x(x), y(y) {}
  Point() : Point(0, 0, 0) {}
};

bool lessID(const Point &p1, const Point &p2) { return p1.id < p2.id; }
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; }

using PointList = vector<Point>;

class KdTree {
  struct Node {
    // int depth;
    Point location;
    unique_ptr<Node> left;
    unique_ptr<Node> right;
  };

  using NodePtr = unique_ptr<Node>;

  NodePtr build(PointList &points, int depth) {
    auto k = points.size();
    if (k < 1)
      return nullptr;

    auto node = make_unique<Node>();
    if (k == 1) {
      node->location = points.front();
      return node;
    }

    if (depth % 2 == 0) {
      // horizontal
      sort(points.begin(), points.end(), lessX);
    } else {
      // vertical
      sort(points.begin(), points.end(), lessY);
    }
    auto mid = points.begin() + (points.size() / 2);
    auto left = PointList(points.begin(), mid);
    auto right = PointList(mid + 1, points.end());
    node->location = *mid;
    node->left = build(left, depth + 1);
    node->right = build(right, depth + 1);
    return node;
  }

public:
  explicit KdTree(PointList &points) : _root(nullptr) {
    _root = build(points, 0);
  }

  PointList find(int sx, int sy, int tx, int ty) {
    PointList points;
    find(*_root, sx, sy, tx, ty, 0, points);
    sort(points.begin(), points.end(), lessID);
    return points;
  }

  void find(const Node &node, int sx, int sy, int tx, int ty, int depth,
            PointList &ans) {
    auto p = node.location;
    if (sx <= p.x && p.x <= tx && sy <= p.y && p.y <= ty) {
      ans.push_back(p);
    }

    if (depth % 2 == 0) {
      // horizontal
      if (node.left != nullptr && sx <= p.x) {
        find(*node.left, sx, sy, tx, ty, depth + 1, ans);
      }
      if (node.right != nullptr && p.x <= tx) {
        find(*node.right, sx, sy, tx, ty, depth + 1, ans);
      }
    } else {
      // vertical
      if (node.left != nullptr && sy <= p.y) {
        find(*node.left, sx, sy, tx, ty, depth + 1, ans);
      }
      if (node.right != nullptr && p.y <= ty) {
        find(*node.right, sx, sy, tx, ty, depth + 1, ans);
      }
    }
  }

private:
  NodePtr _root;
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  int n;
  cin >> n;
  PointList points;
  for (int i = 0; i < n; ++i) {
    int x, y;
    cin >> x >> y;
    points.emplace_back(i, x, y);
  }
  // print(points);
  KdTree kdtree(points);

  int q;
  cin >> q;
  for (int i = 0; i < q; ++i) {
    int sx, sy, tx, ty;
    cin >> sx >> tx >> sy >> ty;
    auto found = kdtree.find(sx, sy, tx, ty);
    for (auto &p : found) {
      cout << p.id << endl;
    }
    cout << endl;
  }

  return 0;
}