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

#include <bits/stdc++.h>
using namespace std;

template <typename T, int B = 3> class persistent_array {
public:
  static const int mask = (1 << B) - 1;
  struct node {
    T val;
    node *ch[1 << B];
    node() { memset(ch, 0, sizeof(node *) << B); }
    node(T v) : val(v) {}
    node(node *t) { memcpy(ch, t->ch, sizeof(node *) << B); }
  };

private:
  const int h, n;
  const T id;
  node *root;
  int get_h(int n) {
    int cnt = 0, tmp = 1;
    while (tmp < n)
      tmp <<= B, cnt++;
    return cnt;
  }
  node *fix(node *t, int p, T val, int x) {
    if (x < 0)
      return new node(val);
    t = t ? new node(t) : new node();
    int to = (p >> x) & mask;
    t->ch[to] = fix(t->ch[to], p, val, x - B);
    return t;
  }
  node *dfs(int p, int d, const vector<T> &v) {
    if ((int)v.size() <= (p << d))
      return nullptr;
    if (d == 0)
      return new node(v[p]);
    node *t = new node;
    for (int i = 0; i < (1 << B); i++) {
      t->ch[i] = dfs((p << B) | i, d - B, v);
    }
    return t;
  }
  node *init(const vector<T> &v) { return dfs(0, h * B, v); }

public:
  persistent_array(int n_, T id_ = 0)
      : h(get_h(n_)), n(1 << (h * B)), id(id_), root(nullptr) {}
  persistent_array(const vector<T> &v, T id_ = 0)
      : h(get_h(v.size)), n(1 << (h * B)), id(id_), root(init(v)) {}
  T set(int p, T val) {
    root = fix(root, p, val, B * (h - 1));
    return val;
  }
  T operator[](int p) const {
    node *t = root;
    for (int i = (h - 1) * B; i >= 0; i -= B) {
      if (!t)
        return id;
      t = t->ch[(p >> i) & mask];
    }
    return t ? t->val : id;
  }
  void trip_version(node *ver) { root = ver; }
  node *get_version() const { return root; }
};

class persistent_union_find {
public:
  using pa = persistent_array<int>;
  using node = typename pa::node;

private:
  pa par;
  int find(int a) { return par[a] < 0 ? a : par.set(a, find(par[a])); }

public:
  persistent_union_find(int n) : par(n, -1) {}
  bool same(int a, int b) { return find(a) == find(b); }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a == b)
      return;
    if (par[a] > par[b])
      swap(a, b);
    par.set(a, par[a] + par[b]);
    par.set(b, a);
  }
  int size(int a) { return -par[find(a)]; }
  void trip_version(node *ver) { par.trip_version(ver); }
  node *get_version() { return par.get_version(); }
};

int main() {
  ios::sync_with_stdio(false), cin.tie(0);
  int n, q;
  cin >> n >> q;
  persistent_union_find uf(n);
  while (q--) {
    int com, x, y;
    cin >> com >> x >> y;
    if (com == 0) {
      uf.unite(x, y);
    } else {
      printf("%d\\n", uf.same(x, y));
    }
  }
  return 0;
}


Error: MLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

template <typename T, int B = 4> class persistent_array {
public:
  static const int mask = (1 << B) - 1;
  struct node {
    T val;
    node *ch[1 << B];
    node() {
      for (int i = 0; i < (1 << B); i++)
        ch[i] = nullptr;
    }
    node(T v) : val(v) {}
    node(node *t) { memcpy(ch, t->ch, sizeof(node *) << B); }
  };

private:
  const int h, n;
  const T id;
  node *root;
  int get_h(int n) {
    int cnt = 0, tmp = 1;
    while (tmp < n)
      tmp <<= B, cnt++;
    return cnt;
  }
  node *fix(node *t, int p, T val, int x) {
    if (x < 0)
      return new node(val);
    t = t ? new node(t) : new node();
    int to = (p >> x) & mask;
    t->ch[to] = fix(t->ch[to], p, val, x - B);
    return t;
  }
  node *dfs(int p, int d, const vector<T> &v) {
    if ((int)v.size() <= (p << d))
      return nullptr;
    if (d == 0)
      return new node(v[p]);
    node *t = new node;
    for (int i = 0; i < (1 << B); i++) {
      t->ch[i] = dfs((p << B) | i, d - B, v);
    }
    return t;
  }
  node *init(const vector<T> &v) { return dfs(0, h * B, v); }

public:
  persistent_array(int n_, T id_ = 0)
      : h(get_h(n_)), n(1 << (h * B)), id(id_), root(nullptr) {}
  persistent_array(const vector<T> &v, T id_ = 0)
      : h(get_h(v.size)), n(1 << (h * B)), id(id_), root(init(v)) {}
  T set(int p, T val) {
    root = fix(root, p, val, B * (h - 1));
    return val;
  }
  T operator[](int p) const {
    node *t = root;
    for (int i = (h - 1) * B; i >= 0; i -= B) {
      if (!t)
        return id;
      t = t->ch[(p >> i) & mask];
    }
    return t ? t->val : id;
  }
  void trip_version(node *t) { root = t; }
  node *get_version() const { return root; }
};

class persistent_union_find {
public:
  using pa = persistent_array<int>;
  using node = typename pa::node;

private:
  pa par;
  int find(int a) { return par[a] < 0 ? a : par.set(a, find(par[a])); }

public:
  persistent_union_find(int n) : par(n, -1) {}
  bool same(int a, int b) { return find(a) == find(b); }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a == b)
      return;
    if (par[a] > par[b])
      swap(a, b);
    par.set(a, par[a] + par[b]);
    par.set(b, a);
  }
  int size(int a) { return -par[find(a)]; }
  void trip_version(node *t) { par.trip_version(t); }
  node *get_version() { return par.get_version(); }
};

int main() {
  ios::sync_with_stdio(false), cin.tie(0);
  int n, q;
  cin >> n >> q;
  persistent_union_find uf(n);
  while (q--) {
    int com, x, y;
    cin >> com >> x >> y;
    if (com == 0) {
      uf.unite(x, y);
    } else {
      printf("%d\\n", uf.same(x, y));
    }
  }
  return 0;
}