Problem p00609 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <bits/stdc++.h>

using namespace std;

#define whole(xs) (xs).begin(), (xs).end()

template <class T> ostream &operator<<(ostream &os, const vector<T> &vs) {
  if (vs.empty())
    return os << "[]";
  os << "[" << vs[0];
  for (int i = 1; i < vs.size(); i++)
    os << " " << vs[i];
  return os << "]";
}

typedef short Long;
// typedef long long Long;

/* 領域木ライブラリ.
 *   Node : 領域木の各ノード
 *   construct(Point[] ps)
 *      : ps(x座標でソートされた点集合)をもとに領域木を作成.
 *   query(Node root, long sx, long sy, long gx, long gy, vector<Point>& result)
 *      : rootに含まれる点のうち, (sx <= x <= gx && sy <= y <=
 * gy)なるもののをresultに入れる
 */

int ans;
struct Point {
  Long x, y;
  Point() {}
  Point(Long x, Long y) : x(x), y(y) {}
  static bool compY(const Point &a, const Point &b) {
    return a.y == b.y ? a.x < b.x : a.y < b.y;
  }
};
Long dist2(const Point &a, const Point &b);
struct Node {
  Node *parent;
  shared_ptr<Node> left, right;
  Long left_x, right_x;
  vector<Point> ys;
};
ostream &operator<<(ostream &os, const Node &node) {
  return os << "( [" << node.left_x << "," << node.right_x << "] -> "
            << node.ys;
}
shared_ptr<Node> construct(vector<Point>::iterator lb,
                           vector<Point>::iterator ub) {
  auto n = make_shared<Node>();
  n->left_x = lb->x;
  n->right_x = (ub - 1)->x;
  if (ub - lb == 1) {
    n->ys.clear();
    n->ys.push_back(*lb);
    return n;
  }
  auto mid = lb + (ub - lb) / 2;
  auto l = construct(lb, mid), r = construct(mid, ub);
  l->parent = n.get();
  r->parent = n.get();
  n->left = l;
  n->right = r;
  size_t li = 0, ri = 0;
  while (true) {
    if (li >= l->ys.size() && ri >= r->ys.size())
      break;
    if (li >= l->ys.size())
      n->ys.push_back(r->ys[ri++]);
    else if (ri >= r->ys.size())
      n->ys.push_back(l->ys[li++]);
    else {
      n->ys.push_back(l->ys[li].y < r->ys[ri].y ? l->ys[li++] : r->ys[ri++]);
    }
  }
  return n;
}

int AN, BN, R;
vector<Point> A, B;

/* (sx <= x <= gx && sy <= y <= gy)なる点(x, y)の個数を返す
 * 閉区間なの注意 */
void query(shared_ptr<Node> &root, Long sx, Long sy, Long gx, Long gy,
           const Point &p) {
  if (gx < root->left_x || root->right_x < sx) {
    // do nothing;
  } else if (sx <= root->left_x && root->right_x <= gx) {
    vector<Point> &ys = root->ys;
    auto lb = upper_bound(whole(ys), Point(0, sy - 1), Point::compY);
    auto ub = lower_bound(whole(ys), Point(0, gy + 1), Point::compY);
    for (auto it = lb; it < ub; it++) {
      if (dist2(*it, p) <= (4 * R) * (4 * R))
        ans++;
    }
    // result.insert(result.end(), lb, ub);
  } else {
    query(root->left, sx, sy, gx, gy, p);
    query(root->right, sx, sy, gx, gy, p);
  }
}

bool input() {
  cin >> AN >> BN >> R;
  if (AN == 0 && BN == 0 && R == 0)
    return false;
  A.clear();
  B.clear();
  A.resize(AN);
  B.resize(BN);
  for (int i = 0; i < AN; i++) {
    cin >> A[i].x >> A[i].y;
  }
  for (int i = 0; i < BN; i++) {
    cin >> B[i].x >> B[i].y;
  }
  return true;
}

ostream &operator<<(ostream &os, const Point &p) {
  return os << "(" << p.x << "," << p.y << ")";
}
bool operator<(const Point &a, const Point &b) {
  return a.x == b.x ? a.y < b.y : a.x < b.x;
}

Long dist2(const Point &a, const Point &b) {
  Long dx = a.x - b.x, dy = a.y - b.y;
  return dx * dx + dy * dy;
}
void solve() {
  sort(whole(B));
  auto rtree = construct(whole(B));
  ans = 0;
  for (int i = 0; i < AN; i++) {
    Point &p = A[i];
    query(rtree, p.x - 4 * R, p.y - 4 * R, p.x + 4 * R, p.y + 4 * R, p);
  }
  cout << ans << endl;
}

int main() {
  while (input())
    solve();
  return 0;
}

Error: MLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define whole(xs) (xs).begin(), (xs).end()

template <class T> ostream &operator<<(ostream &os, const vector<T> &vs) {
  if (vs.empty())
    return os << "[]";
  os << "[" << vs[0];
  for (int i = 1; i < vs.size(); i++)
    os << " " << vs[i];
  return os << "]";
}

typedef short Long;
// typedef long long Long;

/* 領域木ライブラリ.
 *   Node : 領域木の各ノード
 *   construct(Point[] ps)
 *      : ps(x座標でソートされた点集合)をもとに領域木を作成.
 *   query(Node root, long sx, long sy, long gx, long gy, vector<Point>& result)
 *      : rootに含まれる点のうち, (sx <= x <= gx && sy <= y <=
 * gy)なるもののをresultに入れる
 */

int ans;
struct Point {
  Long x, y;
  char type;
  Point() {}
  Point(Long x, Long y, char type = 'c') : x(x), y(y), type(type) {}
  static bool compY(const Point &a, const Point &b) {
    return a.y == b.y ? a.x < b.x : a.y < b.y;
  }
};
Long dist2(const Point &a, const Point &b);
struct Node {
  Node *parent;
  shared_ptr<Node> left, right;
  Long left_x, right_x;
  vector<Point> ys;
};
ostream &operator<<(ostream &os, const Node &node) {
  return os << "( [" << node.left_x << "," << node.right_x << "] -> "
            << node.ys;
}
shared_ptr<Node> construct(vector<Point>::iterator lb,
                           vector<Point>::iterator ub) {
  auto n = make_shared<Node>();
  n->left_x = lb->x;
  n->right_x = (ub - 1)->x;
  if (ub - lb == 1) {
    n->ys.clear();
    n->ys.push_back(*lb);
    return n;
  }
  auto mid = lb + (ub - lb) / 2;
  auto l = construct(lb, mid), r = construct(mid, ub);
  l->parent = n.get();
  r->parent = n.get();
  n->left = l;
  n->right = r;
  size_t li = 0, ri = 0;
  while (true) {
    if (li >= l->ys.size() && ri >= r->ys.size())
      break;
    if (li >= l->ys.size())
      n->ys.push_back(r->ys[ri++]);
    else if (ri >= r->ys.size())
      n->ys.push_back(l->ys[li++]);
    else {
      n->ys.push_back(l->ys[li].y < r->ys[ri].y ? l->ys[li++] : r->ys[ri++]);
    }
  }
  return n;
}

int AN, BN, R;
vector<Point> A, B;

/* (sx <= x <= gx && sy <= y <= gy)なる点(x, y)の個数を返す
 * 閉区間なの注意 */
void query(shared_ptr<Node> &root, Long sx, Long sy, Long gx, Long gy,
           const Point &p) {
  if (gx < root->left_x || root->right_x < sx) {
    // do nothing;
  } else if (sx <= root->left_x && root->right_x <= gx) {
    vector<Point> &ys = root->ys;
    auto lb = upper_bound(whole(ys), Point(0, sy - 1), Point::compY);
    auto ub = lower_bound(whole(ys), Point(0, gy + 1), Point::compY);
    for (auto it = lb; it < ub; it++) {
      if (it->type == 'a')
        continue;
      if (dist2(*it, p) <= (4 * R) * (4 * R))
        ans++;
    }
    // result.insert(result.end(), lb, ub);
  } else {
    query(root->left, sx, sy, gx, gy, p);
    query(root->right, sx, sy, gx, gy, p);
  }
}

bool input() {
  cin >> AN >> BN >> R;
  if (AN == 0 && BN == 0 && R == 0)
    return false;
  A.clear();
  B.clear();
  A.resize(AN);
  B.resize(BN);
  for (int i = 0; i < AN; i++) {
    cin >> A[i].x >> A[i].y;
    A[i].type = 'a';
  }
  for (int i = 0; i < BN; i++) {
    cin >> B[i].x >> B[i].y;
    B[i].type = 'b';
  }
  return true;
}

ostream &operator<<(ostream &os, const Point &p) {
  return os << "(" << p.x << "," << p.y << ")";
}
bool operator<(const Point &a, const Point &b) {
  return a.x == b.x ? a.y < b.y : a.x < b.x;
}

Long dist2(const Point &a, const Point &b) {
  Long dx = a.x - b.x, dy = a.y - b.y;
  return dx * dx + dy * dy;
}
void solve() {
  vector<Point> p = A;
  B.insert(B.end(), A.begin(), A.end());
  sort(whole(B));
  auto rtree = construct(whole(B));
  p.clear();
  ans = 0;
  for (int i = 0; i < AN; i++) {
    Point &p = A[i];
    query(rtree, p.x - 4 * R, p.y - 4 * R, p.x + 4 * R, p.y + 4 * R, p);
  }
  cout << ans << endl;
}

int main() {
  while (input())
    solve();
  return 0;
}