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

#include <bits/stdc++.h>

using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
typedef pair<ll, ll> Pll;
typedef vector<int> Vi;
// typedef tuple<int, int, int> T;
#define FOR(i, s, x) for (int i = s; i < (int)(x); i++)
#define REP(i, x) FOR(i, 0, x)
#define ALL(c) c.begin(), c.end()
#define DUMP(x) cerr << #x << " = " << (x) << endl
#define UNIQUE(c) sort(ALL(c)), c.erase(unique(ALL(c)), c.end())

const int dr[4] = {-1, 0, 1, 0};
const int dc[4] = {0, 1, 0, -1};

template <typename T> struct Data {
  using type = T;
  type data;

  Data(type data) : data(data) {}
  static bool compare(const Data<type> &data1, const Data<type> &data2,
                      size_t base) {
    const size_t dim = data1.data.size();
    for (size_t i = 0; i < dim; i++) {
      if (data1.data[(i + base) % dim] != data2.data[(i + base) % dim]) {
        return data1.data[(i + base) % dim] < data2.data[(i + base) % dim];
      }
    }
    return false;
  }
  bool operator==(const Data<type> &data) const {
    return this->data == data.data;
  }
};

struct Node {
  size_t dim, idx;
  Node *left, *right;
};

template <typename T> struct Region {
  T low, high;
  Region(T low, T high) : low(low), high(high){};
  bool include(T data) const {
    size_t dim = this->low.size();
    for (size_t i = 0; i < dim; i++) {
      if (data[i] < low[i] or high[i] < data[i]) {
        return false;
      }
    }
    return true;
  }
  bool include(Region<T> region) const {
    size_t dim = this->low.size();
    for (size_t bit = 0; bit < (1U << dim); bit++) {
      T data;
      for (size_t i = 0; i < dim; i++) {
        if (bit >> i & 1) {
          data[i] = region.low[i];
        } else {
          data[i] = region.high[i];
        }
      }
      if (not this->include(data)) {
        return false;
      }
    }
    return true;
  }

  bool overlap(Region<T> region) const {
    size_t dim = this->low.size();
    for (size_t bit = 0; bit < (1U << dim); bit++) {
      T data;
      for (size_t i = 0; i < dim; i++) {
        if (bit >> i & 1) {
          data[i] = region.low[i];
        } else {
          data[i] = region.high[i];
        }
      }
      if (this->include(data)) {
        return true;
      }
    }
    return false;
  }

  friend ostream &operator<<(ostream &os, const Region<T> &region) {
    for (size_t i = 0; i < region.low.size(); i++) {
      if (i > 0) {
        os << 'x';
      }
      os << '[' << region.low[i] << ',' << region.high[i] << ']';
    }
    return os;
  }
};

template <typename DataType> struct KDtree {
  using T = typename DataType::type;
  using VT = typename T::value_type;
  size_t dim;
  size_t leaf_dim;
  std::vector<DataType> data;
  VT min_range, max_range;
  Node *root;

  KDtree(VT min_range = std::numeric_limits<VT>::min(),
         VT max_range = std::numeric_limits<VT>::max())
      : dim(std::tuple_size<T>::value), min_range(min_range),
        max_range(max_range) {}

  Node *create_node(size_t dim, size_t idx) {
    Node *node = (Node *)malloc(sizeof(Node));
    if (node == NULL) {
      return NULL;
    }
    node->dim = dim;
    node->idx = idx;
    node->left = node->right = NULL;
    return node;
  }

  Node *build(std::vector<std::vector<size_t>> data, size_t depth) {
    // std::cout << "Build Depth: " << depth << ", Size: ";
    // REP(i, this->dim) std::cout << data[i].size() << (i + 1 == this->dim ?
    // '\\n' : ' ');
    if (data[0].size() == 1) {
      return create_node(leaf_dim, data[0][0]);
    } else {
      std::vector<std::vector<size_t>> left_data(this->dim),
          right_data(this->dim);
      size_t median_idx =
          static_cast<size_t>(static_cast<int>(data[0].size() - 1) / 2);
      size_t median = data[depth % this->dim][median_idx];
      for (size_t i = 0; i < this->dim; i++) {
        for (size_t j : data[i]) {
          if (this->data[j] == this->data[median] or
              DataType::compare(this->data[j], this->data[median],
                                depth % this->dim)) {
            left_data[i].emplace_back(j);
          } else {
            right_data[i].emplace_back(j);
          }
        }
      }
      // std::cout << "Left: " << left_data[0].size() << ", Right: " <<
      // right_data[0].size() << std::endl;
      assert(left_data[0].size() > 0 and right_data[0].size() > 0);
      Node *node = create_node(depth % this->dim, median);
      node->left = build(left_data, depth + 1);
      node->right = build(right_data, depth + 1);
      return node;
    }
  }

  void build_kd_tree(std::vector<DataType> &data) {
    // std::cout << "Start Build..." << std::endl;
    assert(data.size() > 0);

    this->data = data;
    this->leaf_dim = this->dim;

    std::vector<std::vector<size_t>> _data(this->dim,
                                           std::vector<size_t>(data.size()));
    for (size_t i = 0; i < this->dim; i++) {
      std::iota(_data[i].begin(), _data[i].end(), 0);
      sort(_data[i].begin(), _data[i].end(),
           [&](const size_t &a, const size_t &b) {
             return DataType::compare(this->data[a], this->data[b], i);
           });
    }
    this->root = build(_data, 0);
    // std::cout << "End Build" << std::endl;
  }

  void report_subtree(Node *node, std::vector<size_t> &output) const {
    if (node->dim == this->leaf_dim) {
      output.emplace_back(node->idx);
    } else {
      report_subtree(node->left, output);
      report_subtree(node->right, output);
    }
  }

  std::vector<size_t> range_query(Region<T> query_region) {
    T node_low, node_high;
    node_low.fill(this->min_range);
    node_high.fill(this->max_range);
    Region<T> node_region(node_low, node_high);
    std::vector<size_t> output;
    query(this->root, query_region, node_region, output);
    return output;
  }

  void query(Node *node, Region<T> query_region, Region<T> node_region,
             std::vector<size_t> &output) {
    // std::cout << "(Node, Dim) = (" << node->idx << ',' << node->dim << "), "
    // << node_region << ' ' << query_region << std::endl;
    T data = this->data[node->idx].data;
    if (node->dim == this->leaf_dim) {
      if (query_region.include(data)) {
        output.emplace_back(node->idx);
      }
    } else {
      Region<T> left_region = node_region, right_region = node_region;
      Region<T> left_query_region = query_region,
                right_query_region = query_region;

      left_region.high[node->dim] = data[node->dim];
      right_region.low[node->dim] = data[node->dim];
      if (query_region.low[node->dim] <= data[node->dim] and
          data[node->dim] < query_region.high[node->dim]) {
        left_query_region.high[node->dim] = data[node->dim];
      }
      if (data[node->dim] <= query_region.high[node->dim] and
          query_region.low[node->dim] < data[node->dim]) {
        right_query_region.low[node->dim] = data[node->dim];
      }
      // std::cout << left_region << ' ' << right_region << std::endl;
      // std::cout << left_query_region << ' ' << right_query_region <<
      // std::endl;
      if (left_query_region.include(left_region)) {
        report_subtree(node->left, output);
      } else if (left_region.overlap(left_query_region)) {
        query(node->left, left_query_region, left_region, output);
      }

      if (right_query_region.include(right_region)) {
        report_subtree(node->right, output);
      } else if (right_region.overlap(right_query_region)) {
        query(node->right, right_query_region, right_region, output);
      }
    }
  }

  void output() { output_kdtree(this->root); }

  void output_kdtree(Node *node) {
    std::cout << "idx: " << node->idx << ", (";
    for (size_t i = 0; i < this->dim; i++) {
      if (i > 0) {
        std::cout << ',';
      }
      std::cout << this->data[node->idx].data[i];
    }
    std::cout << ") ";
    if (node->dim != this->leaf_dim) {
      std::cout << "(left, right) = (" << node->left->idx << ','
                << node->right->idx << ')' << std::endl;
      output_kdtree(node->left);
      output_kdtree(node->right);
    } else {
      std::cout << std::endl;
    }
  }
};

int main() {
  // use scanf in CodeForces!
  cin.tie(0);
  ios_base::sync_with_stdio(false);

  using DataType = Data<std::array<int, 2>>;
  KDtree<DataType> kdtree;
  std::vector<DataType> points;
  int N;
  cin >> N;
  REP(i, N) {
    int x, y;
    cin >> x >> y;
    DataType::type tmp = {x, y};
    points.emplace_back(DataType(tmp));
  }
  kdtree.build_kd_tree(points);
  std::vector<size_t> output;
  kdtree.report_subtree(kdtree.root, output);
  // for (size_t i : output) std::cout << i << ' ';
  // std::cout << std::endl;
  // kdtree.output();

  ///*
  int Q;
  cin >> Q;
  REP(i, Q) {
    int sx, tx, sy, ty;
    std::cin >> sx >> tx >> sy >> ty;
    DataType::type low = {sx, sy}, high = {tx, ty};
    Region<DataType::type> query_region(low, high);
    std::vector<size_t> out = kdtree.range_query(query_region);
    sort(ALL(out));
    for (size_t i : out)
      printf("%d\\n", i);
    printf("\\n");
  }
  //*/

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
typedef pair<ll, ll> Pll;
typedef vector<int> Vi;
// typedef tuple<int, int, int> T;
#define FOR(i, s, x) for (int i = s; i < (int)(x); i++)
#define REP(i, x) FOR(i, 0, x)
#define ALL(c) c.begin(), c.end()
#define DUMP(x) cerr << #x << " = " << (x) << endl
#define UNIQUE(c) sort(ALL(c)), c.erase(unique(ALL(c)), c.end())

const int dr[4] = {-1, 0, 1, 0};
const int dc[4] = {0, 1, 0, -1};

template <typename T> struct Data {
  using type = T;
  type data;

  Data(type data) : data(data) {}
  static bool compare(const Data<type> &data1, const Data<type> &data2,
                      size_t base) {
    const size_t dim = data1.data.size();
    for (size_t i = 0; i < dim; i++) {
      if (data1.data[(i + base) % dim] != data2.data[(i + base) % dim]) {
        return data1.data[(i + base) % dim] < data2.data[(i + base) % dim];
      }
    }
    return false;
  }
  bool operator==(const Data<type> &data) const {
    return this->data == data.data;
  }
};

struct Node {
  size_t dim, idx;
  Node *left, *right;
};

template <typename T> struct Region {
  T low, high;
  Region(T low, T high) : low(low), high(high){};
  bool include(T data) const {
    size_t dim = this->low.size();
    for (size_t i = 0; i < dim; i++) {
      if (data[i] < low[i] or high[i] < data[i]) {
        return false;
      }
    }
    return true;
  }
  bool include(Region<T> region) const {
    size_t dim = this->low.size();
    for (size_t bit = 0; bit < (1U << dim); bit++) {
      T data;
      for (size_t i = 0; i < dim; i++) {
        if (bit >> i & 1) {
          data[i] = region.low[i];
        } else {
          data[i] = region.high[i];
        }
      }
      if (not this->include(data)) {
        return false;
      }
    }
    return true;
  }

  bool overlap(Region<T> region) const {
    size_t dim = this->low.size();
    for (size_t bit = 0; bit < (1U << dim); bit++) {
      T data;
      for (size_t i = 0; i < dim; i++) {
        if (bit >> i & 1) {
          data[i] = region.low[i];
        } else {
          data[i] = region.high[i];
        }
      }
      if (this->include(data)) {
        return true;
      }
    }
    return false;
  }

  friend ostream &operator<<(ostream &os, const Region<T> &region) {
    for (size_t i = 0; i < region.low.size(); i++) {
      if (i > 0) {
        os << 'x';
      }
      os << '[' << region.low[i] << ',' << region.high[i] << ']';
    }
    return os;
  }
};

template <typename DataType> struct KDtree {
  using T = typename DataType::type;
  using VT = typename T::value_type;
  size_t dim;
  size_t leaf_dim;
  std::vector<DataType> data;
  VT min_range, max_range;
  Node *root;

  KDtree(VT min_range = std::numeric_limits<VT>::min(),
         VT max_range = std::numeric_limits<VT>::max())
      : dim(std::tuple_size<T>::value), min_range(min_range),
        max_range(max_range) {}

  Node *create_node(size_t dim, size_t idx) {
    Node *node = (Node *)malloc(sizeof(Node));
    if (node == NULL) {
      return NULL;
    }
    node->dim = dim;
    node->idx = idx;
    node->left = node->right = NULL;
    return node;
  }

  Node *build(std::vector<std::vector<size_t>> data, size_t depth) {
    // std::cout << "Build Depth: " << depth << ", Size: ";
    // REP(i, this->dim) std::cout << data[i].size() << (i + 1 == this->dim ?
    // '\\n' : ' ');
    if (data[0].size() == 1) {
      return create_node(leaf_dim, data[0][0]);
    } else {
      std::vector<std::vector<size_t>> left_data(this->dim),
          right_data(this->dim);
      size_t median_idx =
          static_cast<size_t>(static_cast<int>(data[0].size() - 1) / 2);
      size_t median = data[depth % this->dim][median_idx];
      for (size_t i = 0; i < this->dim; i++) {
        for (size_t j : data[i]) {
          if (this->data[j] == this->data[median] or
              DataType::compare(this->data[j], this->data[median],
                                depth % this->dim)) {
            left_data[i].emplace_back(j);
          } else {
            right_data[i].emplace_back(j);
          }
        }
      }
      // std::cout << "Left: " << left_data[0].size() << ", Right: " <<
      // right_data[0].size() << std::endl;
      assert(left_data[0].size() > 0 and right_data[0].size() > 0);
      Node *node = create_node(depth % this->dim, median);
      node->left = build(left_data, depth + 1);
      node->right = build(right_data, depth + 1);
      return node;
    }
  }

  void build_kd_tree(std::vector<DataType> &data) {
    // std::cout << "Start Build..." << std::endl;
    assert(data.size() > 0);

    this->data = data;
    this->leaf_dim = this->dim;

    std::vector<std::vector<size_t>> _data(this->dim,
                                           std::vector<size_t>(data.size()));
    for (size_t i = 0; i < this->dim; i++) {
      std::iota(_data[i].begin(), _data[i].end(), 0);
      sort(_data[i].begin(), _data[i].end(),
           [&](const size_t &a, const size_t &b) {
             return DataType::compare(this->data[a], this->data[b], i);
           });
    }
    this->root = build(_data, 0);
    // std::cout << "End Build" << std::endl;
  }

  void report_subtree(Node *node, std::vector<size_t> &output) const {
    if (node->dim == this->leaf_dim) {
      output.emplace_back(node->idx);
    } else {
      report_subtree(node->left, output);
      report_subtree(node->right, output);
    }
  }

  std::vector<size_t> range_query(Region<T> query_region) {
    T node_low, node_high;
    node_low.fill(this->min_range);
    node_high.fill(this->max_range);
    Region<T> node_region(node_low, node_high);
    std::vector<size_t> output;
    query(this->root, query_region, node_region, output);
    return output;
  }

  void query(Node *node, Region<T> query_region, Region<T> node_region,
             std::vector<size_t> &output) {
    // std::cout << "(Node, Dim) = (" << node->idx << ',' << node->dim << "), "
    // << node_region << ' ' << query_region << std::endl;
    T data = this->data[node->idx].data;
    if (node->dim == this->leaf_dim) {
      if (query_region.include(data)) {
        output.emplace_back(node->idx);
      }
    } else {
      Region<T> left_region = node_region, right_region = node_region;
      Region<T> left_query_region = query_region,
                right_query_region = query_region;

      left_region.high[node->dim] = data[node->dim];
      right_region.low[node->dim] = data[node->dim];
      if (query_region.low[node->dim] <= data[node->dim] and
          data[node->dim] < query_region.high[node->dim]) {
        left_query_region.high[node->dim] = data[node->dim];
      }
      if (data[node->dim] <= query_region.high[node->dim] and
          query_region.low[node->dim] < data[node->dim]) {
        right_query_region.low[node->dim] = data[node->dim];
      }
      // std::cout << left_region << ' ' << right_region << std::endl;
      // std::cout << left_query_region << ' ' << right_query_region <<
      // std::endl;
      if (left_query_region.include(left_region)) {
        report_subtree(node->left, output);
      } else if (left_region.overlap(left_query_region)) {
        query(node->left, left_query_region, left_region, output);
      }

      if (right_query_region.include(right_region)) {
        report_subtree(node->right, output);
      } else if (right_region.overlap(right_query_region)) {
        query(node->right, right_query_region, right_region, output);
      }
    }
  }

  void output() { output_kdtree(this->root); }

  void output_kdtree(Node *node) {
    std::cout << "idx: " << node->idx << ", (";
    for (size_t i = 0; i < this->dim; i++) {
      if (i > 0) {
        std::cout << ',';
      }
      std::cout << this->data[node->idx].data[i];
    }
    std::cout << ") ";
    if (node->dim != this->leaf_dim) {
      std::cout << "(left, right) = (" << node->left->idx << ','
                << node->right->idx << ')' << std::endl;
      output_kdtree(node->left);
      output_kdtree(node->right);
    } else {
      std::cout << std::endl;
    }
  }
};

int main() {
  // use scanf in CodeForces!
  cin.tie(0);
  ios_base::sync_with_stdio(false);

  using DataType = Data<std::array<int, 2>>;
  KDtree<DataType> kdtree;
  std::vector<DataType> points;
  int N;
  cin >> N;
  REP(i, N) {
    int x, y;
    cin >> x >> y;
    DataType::type tmp = {x, y};
    points.emplace_back(DataType(tmp));
  }
  kdtree.build_kd_tree(points);
  std::vector<size_t> output;
  kdtree.report_subtree(kdtree.root, output);
  // for (size_t i : output) std::cout << i << ' ';
  // std::cout << std::endl;
  // kdtree.output();

  ///*
  int Q;
  cin >> Q;
  REP(i, Q) {
    int sx, tx, sy, ty;
    std::cin >> sx >> tx >> sy >> ty;
    DataType::type low = {sx, sy}, high = {tx, ty};
    Region<DataType::type> query_region(low, high);
    std::vector<size_t> out = kdtree.range_query(query_region);
    sort(ALL(out));
    for (size_t i : out)
      std::cout << i << std::endl;
    std::cout << std::endl;
  }
  //*/

  return 0;
}