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

#include <bits/stdc++.h>

#define FOR(i, a, b) for (int i = (a); i < ((int)b); ++i)
#define RFOR(i, a) for (int i = (a); i >= 0; --i)
#define FOE(i, a) for (auto i : a)
#define ALL(c) (c).begin(), (c).end()
#define RALL(c) (c).rbegin(), (c).rend()
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define SUM(x) std::accumulate(ALL(x), 0LL)
#define MIN(v) *std::min_element(v.begin(), v.end())
#define MAX(v) *std::max_element(v.begin(), v.end())
#define EXIST(v, x) (std::find(v.begin(), v.end(), x) != v.end())
#define BIT_ON(bit, i) (((bit >> i) & 1) > 0)

typedef long long LL;
template <typename T> std::vector<T> make_v(size_t a) {
  return std::vector<T>(a);
}
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return std::vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
} // C++14
template <typename T, typename V>
typename std::enable_if<std::is_class<T>::value == 0>::type fill_v(T &t,
                                                                   const V &v) {
  t = v;
}
template <typename T, typename V>
typename std::enable_if<std::is_class<T>::value != 0>::type fill_v(T &t,
                                                                   const V &v) {
  for (auto &e : t)
    fill_v(e, v);
}
template <class T> inline T ceil(T a, T b) { return (a + b - 1) / b; }
template <class T> inline void print(T x) { std::cout << x << std::endl; }
template <class T> inline void print_vec(const std::vector<T> &v) {
  for (int i = 0; i < v.size(); ++i) {
    if (i != 0) {
      std::cout << " ";
    }
    std::cout << v[i];
  }
  std::cout << "\\n";
}
template <class T> inline bool inside(T y, T x, T H, T W) {
  return 0 <= y and y < H and 0 <= x and x < W;
}
template <class T> inline double euclidean_distance(T y1, T x1, T y2, T x2) {
  return sqrt((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2));
}
template <class T> inline double manhattan_distance(T y1, T x1, T y2, T x2) {
  return abs(x1 - x2) + abs(y1 - y2);
}
template <typename T> T &chmin(T &a, const T &b) { return a = std::min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = std::max(a, b); }
template <class T> inline std::vector<T> unique(std::vector<T> v) {
  sort(v.begin(), v.end());
  v.erase(unique(v.begin(), v.end()), v.end());
  return v;
}
// 初項s, 交差dのn個の数列の和
long long sum_of_arithmetic_progression(long long s, long long d, long long n) {
  return n * (2 * s + (n - 1) * d) / 2;
}

const int INF = 1 << 30;
const double EPS = 1e-9;
const long double PI = acos(-1.0);
const std::vector<int> dy4 = {0, -1, 0, 1}, dx4 = {1, 0, -1, 0};
const std::vector<int> dy8 = {0, -1, 0, 1, 1, -1, -1, 1},
                       dx8 = {1, 0, -1, 0, 1, 1, -1, -1};

// 木上のクエリを処理
// ・作成(O(n log n))
// ・ノード間の距離(O(log n))
// ・直径・中心(O(log n))
// ・最小共通素性(O(log n))
// ・ノード間の最大の重み(O(log n))
class QueryOnTree {
private:
  std::unordered_map<int, std::vector<std::pair<int, int>>> tree;
  int log_v = 0;
  std::vector<std::unordered_map<int, int>> parent; // 2^k個上の親
  //    std::vector<std::unordered_map<int, int>> max_weight;    //
  //    2^k個上の親までにでてくる最大の重み
  std::unordered_map<int, int> depth; // 各頂点についての根からの深さ
  //    std::unordered_map<int, long long> distance_from_root;   // 根からの距離
  bool done = false;

public:
  QueryOnTree() {}

  // 無向グラフの場合はu->vとv->uの両方をいれること
  QueryOnTree(
      const std::unordered_map<int, std::vector<std::pair<int, int>>> &tree)
      : tree(tree) {}

  // 無向グラフの場合はu->vとv->uの両方をいれること
  void add_edge(int u, int v, int w) {
    assert(not done);
    tree[u].emplace_back(std::make_pair(v, w));
  }

  void build(int root) {
    this->done = true;

    int n = tree.size();
    if (n == 0) {
      return;
    }

    this->log_v = int(log2(n)) + 1;
    this->parent = std::vector<std::unordered_map<int, int>>(log_v);
    //        this->max_weight = std::vector<std::unordered_map<int,
    //        int>>(log_v); this->distance_from_root[root] = 0;

    dfs(root, -1, 0, 0);

    for (int k = 0; k + 1 < this->log_v; k++) {
      for (const auto &p : tree) {
        int u = p.first;
        if (parent[k][u] < 0) {
          parent[k + 1][u] = -1;
        } else {
          parent[k + 1][u] = parent
              [k]
              [parent
                   [k]
                   [u]]; // uの2^k個上のノードの2^k上のノードはuの2^(k+1)個上のノード
          if (parent[k + 1][u] >= 0) {
            //                        max_weight[k + 1][u] =
            //                        std::max(max_weight[k][u],
            //                        max_weight[k][parent[k][u]]);
          }
        }
      }
    }
  }

  // uとvの距離
  //    long long distance(int u, int v) {
  //        return distance_from_root[u] + distance_from_root[v] - 2 *
  //        distance_from_root[lca(u, v)];
  //    }

  // 木の直径，中心
  //    std::pair<int, long long> center_diameter() {
  //        // node 0から一番遠いnode
  //        int u = 0;
  //        long long max_distance = -1;
  //        for (const auto &p : tree) {
  //            auto node = p.first;
  //            auto dist = this->distance(0, node);
  //            if (dist > max_distance) {
  //                max_distance = dist;
  //                u = node;
  //            }
  //        }
  //
  //        // node uから一番遠いnode
  //        std::unordered_map<long long, int> distance_node;
  //        long long diameter = 0;
  //        for (const auto &p : tree) {
  //            auto node = p.first;
  //            auto dist = this->distance(u, node);
  //            distance_node[dist] = node;
  //            if (dist > diameter) {
  //                diameter = dist;
  //            }
  //        }
  //        assert(diameter != -1);
  //
  //        int center = distance_node[max_distance / 2];
  //
  //        return std::make_pair(center, diameter);
  //    }

  // uとvの間で出現する最大の重み
  //    int maximum_weight(int u, int v) {
  //        int lca = this->lca(u, v);
  //        return std::max(this->maximum_weight_ancestor(u, lca),
  //        this->maximum_weight_ancestor(v, lca));
  //    }

  // uとvの最近共通祖先(O(log n))
  int lca(int u, int v) {
    if (depth[u] > depth[v]) {
      std::swap(u, v);
    }
    for (int k = 0; k < this->log_v; ++k) {
      if ((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }

    if (u == v) {
      return u;
    }

    for (int k = this->log_v - 1; k >= 0; --k) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }

private:
  // 各頂点について，1つ上の親，1つ上の重み，根からの深さ，根からの距離を求める
  void dfs(int u, int p, int d, int w) {
    parent[0][u] = p;
    //        max_weight[0][u] = w;
    depth[u] = d;

    for (const auto &a : tree[u]) {
      int v = a.first;
      int dist = a.second;
      if (v != p) {
        //                distance_from_root[v] = distance_from_root[u] + dist;
        dfs(v, u, d + 1, dist);
      }
    }
  }

  // uとuの祖先の間の最大の重みを取得
  //    int maximum_weight_ancestor(int u, int ancestor) {
  //        int res = -INF;
  //        int d = depth[u] - depth[ancestor];
  //        for (int k = 0; k < log_v; k ++) {
  //            if ((d >> k) & 1) {
  //                res = std::max(res, max_weight[k][u]);
  //                u = parent[k][u];
  //            }
  //        }
  //        return res;
  //    }
};

using namespace std;

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

  int N;
  cin >> N;

  QueryOnTree qut;
  FOR(i, 0, N) {
    int K;
    cin >> K;
    FOR(j, 0, K) {
      int C;
      cin >> C;
      qut.add_edge(i, C, 1);
      qut.add_edge(C, i, 1);
    }
  }
  qut.build(0);

  int Q;
  cin >> Q;
  FOR(i, 0, Q) {
    int U, V;
    cin >> U >> V;
    print(qut.lca(U, V));
  }

  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>

#define FOR(i, a, b) for (int i = (a); i < ((int)b); ++i)
#define RFOR(i, a) for (int i = (a); i >= 0; --i)
#define FOE(i, a) for (auto i : a)
#define ALL(c) (c).begin(), (c).end()
#define RALL(c) (c).rbegin(), (c).rend()
#define DUMP(x) cerr << #x << " = " << (x) << endl;
#define SUM(x) std::accumulate(ALL(x), 0LL)
#define MIN(v) *std::min_element(v.begin(), v.end())
#define MAX(v) *std::max_element(v.begin(), v.end())
#define EXIST(v, x) (std::find(v.begin(), v.end(), x) != v.end())
#define BIT_ON(bit, i) (((bit >> i) & 1) > 0)

typedef long long LL;
template <typename T> std::vector<T> make_v(size_t a) {
  return std::vector<T>(a);
}
template <typename T, typename... Ts> auto make_v(size_t a, Ts... ts) {
  return std::vector<decltype(make_v<T>(ts...))>(a, make_v<T>(ts...));
} // C++14
template <typename T, typename V>
typename std::enable_if<std::is_class<T>::value == 0>::type fill_v(T &t,
                                                                   const V &v) {
  t = v;
}
template <typename T, typename V>
typename std::enable_if<std::is_class<T>::value != 0>::type fill_v(T &t,
                                                                   const V &v) {
  for (auto &e : t)
    fill_v(e, v);
}
template <class T> inline T ceil(T a, T b) { return (a + b - 1) / b; }
template <class T> inline void print(T x) { std::cout << x << std::endl; }
template <class T> inline void print_vec(const std::vector<T> &v) {
  for (int i = 0; i < v.size(); ++i) {
    if (i != 0) {
      std::cout << " ";
    }
    std::cout << v[i];
  }
  std::cout << "\\n";
}
template <class T> inline bool inside(T y, T x, T H, T W) {
  return 0 <= y and y < H and 0 <= x and x < W;
}
template <class T> inline double euclidean_distance(T y1, T x1, T y2, T x2) {
  return sqrt((x1 - x2) * (x1 - x2) + (y1 - y2) * (y1 - y2));
}
template <class T> inline double manhattan_distance(T y1, T x1, T y2, T x2) {
  return abs(x1 - x2) + abs(y1 - y2);
}
template <typename T> T &chmin(T &a, const T &b) { return a = std::min(a, b); }
template <typename T> T &chmax(T &a, const T &b) { return a = std::max(a, b); }
template <class T> inline std::vector<T> unique(std::vector<T> v) {
  sort(v.begin(), v.end());
  v.erase(unique(v.begin(), v.end()), v.end());
  return v;
}
// 初項s, 交差dのn個の数列の和
long long sum_of_arithmetic_progression(long long s, long long d, long long n) {
  return n * (2 * s + (n - 1) * d) / 2;
}

const int INF = 1 << 30;
const double EPS = 1e-9;
const long double PI = acos(-1.0);
const std::vector<int> dy4 = {0, -1, 0, 1}, dx4 = {1, 0, -1, 0};
const std::vector<int> dy8 = {0, -1, 0, 1, 1, -1, -1, 1},
                       dx8 = {1, 0, -1, 0, 1, 1, -1, -1};

// 木上のクエリを処理
// ・作成(O(n log n))
// ・ノード間の距離(O(log n))
// ・直径・中心(O(log n))
// ・最小共通素性(O(log n))
// ・ノード間の最大の重み(O(log n))
class QueryOnTree {
private:
  std::unordered_map<int, std::vector<std::pair<int, int>>> tree;
  int log_v = 0;
  std::vector<std::unordered_map<int, int>> parent; // 2^k個上の親
  std::vector<std::unordered_map<int, int>>
      max_weight; // 2^k個上の親までにでてくる最大の重み
  std::unordered_map<int, int> depth; // 各頂点についての根からの深さ
  std::unordered_map<int, long long> distance_from_root; // 根からの距離

public:
  QueryOnTree() {}

  // 無向グラフの場合はu->vとv->uの両方をいれること
  QueryOnTree(
      const std::unordered_map<int, std::vector<std::pair<int, int>>> &tree)
      : tree(tree) {}

  // 無向グラフの場合はu->vとv->uの両方をいれること
  void add_edge(int u, int v, int w) {
    tree[u].emplace_back(std::make_pair(v, w));
  }

  void build(int root) {
    int n = tree.size();
    if (n == 0) {
      return;
    }

    this->log_v = int(log2(n)) + 1;
    this->parent = std::vector<std::unordered_map<int, int>>(log_v);
    this->max_weight = std::vector<std::unordered_map<int, int>>(log_v);
    this->distance_from_root[root] = 0;

    dfs(root, -1, 0, 0);

    for (int k = 0; k + 1 < this->log_v; k++) {
      for (const auto &p : tree) {
        int u = p.first;
        if (parent[k][u] < 0) {
          parent[k + 1][u] = -1;
        } else {
          parent[k + 1][u] = parent
              [k]
              [parent
                   [k]
                   [u]]; // uの2^k個上のノードの2^k上のノードはuの2^(k+1)個上のノード
          if (parent[k + 1][u] >= 0) {
            max_weight[k + 1][u] =
                std::max(max_weight[k][u], max_weight[k][parent[k][u]]);
          }
        }
      }
    }
  }

  // uとvの距離
  long long distance(int u, int v) {
    return distance_from_root[u] + distance_from_root[v] -
           2 * distance_from_root[lca(u, v)];
  }

  // 木の直径，中心
  std::pair<int, long long> center_diameter() {
    // node 0から一番遠いnode
    int u = 0;
    long long max_distance = -1;
    for (const auto &p : tree) {
      auto node = p.first;
      auto dist = this->distance(0, node);
      if (dist > max_distance) {
        max_distance = dist;
        u = node;
      }
    }

    // node uから一番遠いnode
    std::unordered_map<long long, int> distance_node;
    long long diameter = 0;
    for (const auto &p : tree) {
      auto node = p.first;
      auto dist = this->distance(u, node);
      distance_node[dist] = node;
      if (dist > diameter) {
        diameter = dist;
      }
    }
    assert(diameter != -1);

    int center = distance_node[max_distance / 2];

    return std::make_pair(center, diameter);
  }

  // uとvの間で出現する最大の重み
  int maximum_weight(int u, int v) {
    int lca = this->lca(u, v);
    return std::max(this->maximum_weight_ancestor(u, lca),
                    this->maximum_weight_ancestor(v, lca));
  }

  // uとvの最近共通祖先(O(log n))
  int lca(int u, int v) {
    if (depth[u] > depth[v]) {
      std::swap(u, v);
    }
    for (int k = 0; k < this->log_v; ++k) {
      if ((depth[v] - depth[u]) >> k & 1) {
        v = parent[k][v];
      }
    }

    if (u == v) {
      return u;
    }

    for (int k = this->log_v - 1; k >= 0; --k) {
      if (parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }
    }
    return parent[0][u];
  }

private:
  // 各頂点について，1つ上の親，1つ上の重み，根からの深さ，根からの距離を求める
  void dfs(int u, int p, int d, int w) {
    parent[0][u] = p;
    max_weight[0][u] = w;
    depth[u] = d;

    for (const auto &a : tree[u]) {
      int v = a.first;
      int dist = a.second;
      if (v != p) {
        distance_from_root[v] = distance_from_root[u] + dist;
        dfs(v, u, d + 1, dist);
      }
    }
  }

  // uとuの祖先の間の最大の重みを取得
  int maximum_weight_ancestor(int u, int ancestor) {
    int res = -INF;
    int d = depth[u] - depth[ancestor];
    for (int k = 0; k < log_v; k++) {
      if ((d >> k) & 1) {
        res = std::max(res, max_weight[k][u]);
        u = parent[k][u];
      }
    }
    return res;
  }
};

using namespace std;

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

  int N;
  cin >> N;

  QueryOnTree qut;
  FOR(i, 0, N) {
    int K;
    cin >> K;
    FOR(j, 0, K) {
      int C;
      cin >> C;
      qut.add_edge(i, C, 1);
      qut.add_edge(C, i, 1);
    }
  }
  qut.build(0);

  int Q;
  cin >> Q;
  FOR(i, 0, Q) {
    int U, V;
    cin >> U >> V;
    print(qut.lca(U, V));
  }

  return 0;
}