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

#include <algorithm>
#include <cstdio>
#include <queue>
#include <utility>
#include <vector>

using namespace std;

template <class T> using lp_queue = priority_queue<T, deque<T>, greater<T>>;

struct Point {
  int first, second;
  size_t index;
  Point(int f = 0, int s = 0, size_t i = 0) : first(f), second(s), index(i) {}
  bool operator<(const Point &rhs) const {
    if (first != rhs.first) {
      return first < rhs.first;
    } else {
      return second < rhs.second;
    }
  }
};

int main() {
  size_t n;
  scanf("%zu", &n);

  vector<Point> S(n), T(n);
  for (size_t i = 0; i < n; ++i) {
    int x, y;
    scanf("%d %d", &x, &y);

    S[i].first = T[i].second = x;
    S[i].second = T[i].first = y;
    S[i].index = T[i].index = i;
  }

  sort(S.begin(), S.end());
  sort(T.begin(), T.end());

  size_t q;
  scanf("%zu", &q);

  for (size_t i = 0; i < q; ++i) {
    int sx, tx, sy, ty;
    scanf("%d %d %d %d", &sx, &tx, &sy, &ty);

    lp_queue<size_t> q;
    size_t x_lb = lower_bound(S.begin(), S.end(), Point(sx, sy)) - S.begin();
    size_t x_ub = upper_bound(S.begin(), S.end(), Point(tx, ty)) - S.begin();
    size_t y_lb = lower_bound(T.begin(), T.end(), Point(sy, sx)) - T.begin();
    size_t y_ub = upper_bound(T.begin(), T.end(), Point(ty, tx)) - T.begin();

    ptrdiff_t x_diff = x_ub - x_lb, y_diff = y_ub - y_lb;
    if (x_diff < y_diff) {
      for (size_t i = x_lb; i < x_ub; ++i) {
        int y = S[i].second;
        if (sy <= y && y <= ty)
          q.push(S[i].index);
      }
    } else {
      for (size_t i = y_lb; i < y_ub; ++i) {
        int x = T[i].second;
        if (sx <= x && x <= tx)
          q.push(T[i].index);
      }
    }

    while (!q.empty()) {
      printf("%zd\\n", q.top());
      q.pop();
    }
    printf("\\n");
  }

  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <set>
#include <utility>
#include <vector>

using namespace std;

struct KDimTree {
  using Point = pair<int, int>;
  struct Node {
    size_t index;
    Point p;
    Node *l, *r;
    Node(const Point &p, size_t i = 0) : index(i), p(p), l(NULL), r(NULL) {}
  } *root;
  struct Comp {
    bool operator()(int d, const Point &p, const Point &q) {
      return d ? p.first < q.first : p.second < q.second;
    }
  } comp;
  KDimTree() : root(NULL) {}
  void insert(const Point &p, size_t i = 0) { root = insert(root, 0, p, i); }
  Node *insert(Node *t, int d, const Point &p, size_t i = 0) {
    if (t == NULL)
      return new Node(p, i);
    if (comp(d, p, t->p)) {
      t->l = insert(t->l, !d, p, i);
    } else {
      t->r = insert(t->r, !d, p, i);
    };
    return t;
  }
  template <class Output>
  bool search(const Point &ld, const Point &ru, Output &out) {
    search(root, 0, ld, ru, out);
    return !out.empty();
  }
  template <class Output>
  void search(Node *t, int d, const Point &ld, const Point &ru, Output &out) {
    if (t == NULL)
      return;

    const Point &p = t->p;
    if (ld.first <= p.first && p.first <= ru.first)
      if (ld.second <= p.second && p.second <= ru.second)
        out.insert(t);

    if (!comp(d, p, ld))
      search(t->l, !d, ld, ru, out);
    if (!comp(d, ru, p))
      search(t->r, !d, ld, ru, out);
  }
};

int main() {
  size_t n;
  scanf("%zu", &n);

  KDimTree tree;
  for (size_t i = 0; i < n; ++i) {
    KDimTree::Point p;
    scanf("%d %d", &p.first, &p.second);

    tree.insert(p, i);
  }

  size_t q;
  scanf("%zu", &q);

  struct IndexCompare {
    bool operator()(const KDimTree::Node *x, const KDimTree::Node *y) {
      return x->index < y->index;
    }
  };

  for (size_t i = 0; i < q; ++i) {
    KDimTree::Point ld, ru;
    scanf("%d %d %d %d", &ld.first, &ru.first, &ld.second, &ru.second);

    set<KDimTree::Node *> out;
    tree.search(ld, ru, out);

    for (const KDimTree::Node *x : out)
      printf("%zd\\n", x->index);

    printf("\\n");
  }

  return 0;
}