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

#include <bits/stdc++.h>
#define MOD 1000000007
// #define MOD 998244353

class mint {
public:
  int i;

  mint() : i(0) {}

  template <class T> mint(T x) {
    i = int(x % MOD);
    if (i < 0)
      i += MOD;
  }

  mint operator+(mint x) { return i + x.i; }

  mint operator-(mint x) { return i - x.i; }

  mint operator*(mint x) { return (long long)i * x.i; }

  mint operator/(mint x) { return (long long)i * x.pow(MOD - 2).i; }

  mint inv() { return pow(MOD - 2); }

  template <class T> mint pow(T p) {
    long long r = 1;
    long long t = i;
    for (; p > 0; p >>= 1) {
      if (p & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  template <class T1, class T2> static mint pow(T1 a, T2 b) {
    long long r = 1;
    long long t = (long long)(a % MOD);
    for (; b > 0; b >>= 1) {
      if (b & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  mint &operator+=(mint x) {
    i = (i + x.i) % MOD;
    return *this;
  }

  mint &operator-=(mint x) {
    i = i - x.i;
    if (i < 0)
      i += MOD;
    return *this;
  }

  mint &operator*=(mint x) {
    i = (int)((long long)i * x.i % MOD);
    return *this;
  }

  mint &operator/=(mint x) {
    i = (long long)i * x.pow(MOD - 2).i % MOD;
    return *this;
  }
};

std::ostream &operator<<(std::ostream &os, const mint &m) { return os << m.i; }

struct combination {
  std::vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    assert(n < MOD);
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
} mcomb(200002);

// T has to implement 'operator+' and 'add_parent' method
template <class T> class rtdp {
  class Node {
  public:
    int k;                            // k_th node
    rtdp &nodes;                      // owner object holding all nodes
    const std::vector<int> &children; // children indices
    const int csize;
    std::vector<T> dp; // stores dps of children
    std::vector<T> lmerged;
    std::vector<T> rmerged;

    Node(rtdp &ns, int i)
        : k(i), nodes(ns), children(ns.v[i]), csize(children.size()), dp(csize),
          lmerged(csize + 1), rmerged(csize + 1) {}

    T fc(int parent = -1) {
      const int csize = children.size();
      T m;
      for (int i = 0; i < csize; i++) {
        int ci = children[i];
        if (ci == parent)
          continue;
        dp[i] = nodes.ns[ci].fc(k);
        m = m + dp[i];
      }
      return m.add_parent(k);
    }

    void fp(int parent = -1) {
      const int csize = children.size();
      // assume that dp is already filled in
      // left merge
      for (int i = 0; i < csize; i++) {
        lmerged[i + 1] = lmerged[i] + dp[i];
      }
      // right merge
      for (int i = csize; i > 0; i--) {
        rmerged[i - 1] = rmerged[i] + dp[i - 1];
      }
      for (int i = 0; i < csize; i++) {
        int ci = children[i];
        if (ci == parent)
          continue;
        Node &cnode = nodes.ns[ci];
        for (int j = 0; j < cnode.csize; j++) {
          if (cnode.children[j] == k) {
            cnode.dp[j] = (lmerged[i] + rmerged[i + 1]).add_parent(k);
            break;
          }
        }
        cnode.fp(k);
      }
      return;
    }
  };

  std::vector<std::vector<int>> &v;
  int size;
  std::vector<Node> ns;
  std::vector<T> root_dp;

public:
  rtdp(std::vector<std::vector<int>> &v) : v(v), size(v.size()), root_dp(size) {
    for (int i = 0; i < size; i++) {
      ns.push_back(Node(*this, i));
    }
  }

  void exec_dp(int root = 0) {
    fill_in_children(root);
    fill_in_parent(root);
  }

  void fill_in_children(int root) { ns[root].fc(); }

  void fill_in_parent(int root) { ns[root].fp(); }

  T get_root_dp(int i) {
    Node &r = ns[i];
    return r.lmerged[r.csize].add_parent(i);
  }
};

struct DP {
  int tsize;
  mint count;

  DP(int s = 0, mint c = 1) : tsize(s), count(c) {}

  DP operator+(DP d) {
    int tot_tsize = tsize + d.tsize;
    mint tot_count = count * d.count * mcomb(tot_tsize, d.tsize);
    return DP(tot_tsize, tot_count);
  }

  DP add_parent(int k) { return DP(tsize + 1, count); }
};

int main() {
  int n;
  std::cin >> n;
  std::vector<std::vector<int>> v(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    std::cin >> a >> b;
    a--;
    b--;
    v[a].push_back(b);
    v[b].push_back(a);
  }
  rtdp<DP> r(v);
  r.exec_dp();
  for (int i = 0; i < n; i++) {
    std::cout << r.get_root_dp(i).count << std::endl;
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define MOD 1000000007
// #define MOD 998244353

class mint {
public:
  int i;

  mint() : i(0) {}

  template <class T> mint(T x) {
    i = int(x % MOD);
    if (i < 0)
      i += MOD;
  }

  mint operator+(mint x) { return i + x.i; }

  mint operator-(mint x) { return i - x.i; }

  mint operator*(mint x) { return (long long)i * x.i; }

  mint operator/(mint x) { return (long long)i * x.pow(MOD - 2).i; }

  mint inv() { return pow(MOD - 2); }

  template <class T> mint pow(T p) {
    long long r = 1;
    long long t = i;
    for (; p > 0; p >>= 1) {
      if (p & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  template <class T1, class T2> static mint pow(T1 a, T2 b) {
    long long r = 1;
    long long t = (long long)(a % MOD);
    for (; b > 0; b >>= 1) {
      if (b & 1)
        r = r * t % MOD;
      t = t * t % MOD;
    }
    return r;
  }

  mint &operator+=(mint x) {
    i = (i + x.i) % MOD;
    return *this;
  }

  mint &operator-=(mint x) {
    i = i - x.i;
    if (i < 0)
      i += MOD;
    return *this;
  }

  mint &operator*=(mint x) {
    i = (int)((long long)i * x.i % MOD);
    return *this;
  }

  mint &operator/=(mint x) {
    i = (long long)i * x.pow(MOD - 2).i % MOD;
    return *this;
  }
};

std::ostream &operator<<(std::ostream &os, const mint &m) { return os << m.i; }

mint mperm(int i, int j) {
  mint res = 1;
  for (; j > 0; j--) {
    res *= i;
    i--;
  }
  return res;
}

mint mcomb(int i, int j) {
  int k = i - j;
  if (k < j)
    j = k;
  return mperm(i, j) / mperm(j, j);
}

// T has to implement 'operator+' and 'add_parent' method
template <class T> class rtdp {
  class Node {
  public:
    int k;                            // k_th node
    rtdp &nodes;                      // owner object holding all nodes
    const std::vector<int> &children; // children indices
    const int csize;
    std::vector<T> dp; // stores dps of children
    std::vector<T> lmerged;
    std::vector<T> rmerged;

    Node(rtdp &ns, int i)
        : k(i), nodes(ns), children(ns.v[i]), csize(children.size()), dp(csize),
          lmerged(csize + 1), rmerged(csize + 1) {}

    T fc(int parent = -1) {
      const int csize = children.size();
      T m;
      for (int i = 0; i < csize; i++) {
        int ci = children[i];
        if (ci == parent)
          continue;
        dp[i] = nodes.ns[ci].fc(k);
        m = m + dp[i];
      }
      return m.add_parent(k);
    }

    void fp(int parent = -1) {
      const int csize = children.size();
      // assume that dp is already filled in
      // left merge
      for (int i = 0; i < csize; i++) {
        lmerged[i + 1] = lmerged[i] + dp[i];
      }
      // right merge
      for (int i = csize; i > 0; i--) {
        rmerged[i - 1] = rmerged[i] + dp[i - 1];
      }
      for (int i = 0; i < csize; i++) {
        int ci = children[i];
        if (ci == parent)
          continue;
        Node &cnode = nodes.ns[ci];
        for (int j = 0; j < cnode.csize; j++) {
          if (cnode.children[j] == k) {
            cnode.dp[j] = (lmerged[i] + rmerged[i + 1]).add_parent(k);
            break;
          }
        }
        cnode.fp(k);
      }
      return;
    }
  };

  std::vector<std::vector<int>> &v;
  int size;
  std::vector<Node> ns;
  std::vector<T> root_dp;

public:
  rtdp(std::vector<std::vector<int>> &v) : v(v), size(v.size()), root_dp(size) {
    for (int i = 0; i < size; i++) {
      ns.push_back(Node(*this, i));
    }
  }

  void exec_dp(int root = 0) {
    fill_in_children(root);
    fill_in_parent(root);
  }

  void fill_in_children(int root) { ns[root].fc(); }

  void fill_in_parent(int root) { ns[root].fp(); }

  T get_root_dp(int i) {
    Node &r = ns[i];
    return r.lmerged[r.csize].add_parent(i);
  }
};

struct DP {
  int tsize;
  mint count;

  DP(int s = 0, mint c = 1) : tsize(s), count(c) {}

  DP operator+(DP d) {
    int tot_tsize = tsize + d.tsize;
    mint tot_count = count * d.count * mcomb(tot_tsize, d.tsize);
    return DP(tot_tsize, tot_count);
  }

  DP add_parent(int k) { return DP(tsize + 1, count); }
};

int main() {
  int n;
  std::cin >> n;
  std::vector<std::vector<int>> v(n);
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    std::cin >> a >> b;
    a--;
    b--;
    v[a].push_back(b);
    v[b].push_back(a);
  }
  rtdp<DP> r(v);
  r.exec_dp();
  for (int i = 0; i < n; i++) {
    std::cout << r.get_root_dp(i).count << std::endl;
  }
}