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

#ifdef BUILT_LOCAL

#include "niklib.h"

#endif

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
constexpr ll mod = 1000000007;

typedef long long num;

#define bs(n) (1ll << (ll)(n))
#define rep(i, n) for (ll(i) = 0; (i) < (n); (i)++)
#define rrep(i, n) for (ll(i) = 1; (i) <= (n); (i)++)

void extgcd(num a, num b, num &ra, num &rb) {
  if (b == 0) {
    ra = a;
    rb = 0;
  } else {
    num ta, tb;
    extgcd(b, a % b, ta, tb);
    ra = tb;
    rb = ta - tb * (a / b);
  }
}

num inv_mod(num a) {
  num ta, tb;
  extgcd(a, mod, ta, tb);
  return ((ta % mod) + mod) % mod;
}

num dp_factor[300000];

num factor(num n) {
  if (n >= 300000 || n < 0)
    throw;
  if (n == 0)
    return 1;
  if (dp_factor[n] == 0)
    dp_factor[n] = (factor(n - 1) * n) % mod;
  return dp_factor[n];
}

num dp_factor_inv[300000];

num factor_inv(num n) {
  if (n >= 300000 || n < 0)
    throw;
  if (n == 0)
    return 1;
  if (dp_factor_inv[n] == 0)
    dp_factor_inv[n] = (factor_inv(n - 1) * inv_mod(n)) % mod;
  return dp_factor_inv[n];
}

num N, a, b;
vector<set<num>> path;

vector<map<num, num>> dp_count_node;

num count_node_subtree(num node, num parent) {
  if (dp_count_node[node].find(parent) == dp_count_node[node].end()) {
    if (dp_count_node[node].empty()) {
      num r = 1;
      for (auto &e : path[node]) {
        if (e == parent)
          continue;
        r += count_node_subtree(e, node);
      }
      dp_count_node[node][parent] = r;
    } else {
      num r_parent = dp_count_node[node].begin()->first;
      num r_count = dp_count_node[node].begin()->second;
      num r = r_count;
      if (r_parent != 0) {
        r += count_node_subtree(r_parent, node);
      }
      if (parent != 0) {
        r -= count_node_subtree(parent, node);
      }
      dp_count_node[node][parent] = r;
    }
  }
  return dp_count_node[node][parent];
}

vector<map<num, num>> dp_count_dist;

num count_dist_subtree(num node, num parent) {
  if (dp_count_dist[node].find(parent) == dp_count_dist[node].end()) {
    if (dp_count_dist[node].empty()) {
      num r = factor(count_node_subtree(node, parent) - 1);
      for (auto &e : path[node]) {
        if (e == parent)
          continue;
        r = (r * count_dist_subtree(e, node)) % mod;
        r = (r * factor_inv(count_node_subtree(e, node))) % mod;
      }
      dp_count_dist[node][parent] = r;
    } else {
      num r_parent = dp_count_dist[node].begin()->first;
      num r_count = dp_count_dist[node].begin()->second;
      num r = r_count;
      r = (r * (factor_inv(count_node_subtree(node, r_parent) - 1))) % mod;
      r = (r * (factor(count_node_subtree(node, parent) - 1))) % mod;
      if (r_parent != 0) {
        r = (r * factor_inv(count_node_subtree(r_parent, node))) % mod;
        r = (r * count_dist_subtree(r_parent, node)) % mod;
      }
      if (parent != 0) {
        r = (r * factor(count_node_subtree(parent, node))) % mod;
        r = (r * inv_mod(count_dist_subtree(parent, node))) % mod;
      }
      dp_count_dist[node][parent] = r;
    }
  }
  return dp_count_dist[node][parent];
}

int main() {
  cin >> N;
  path.resize(N + 1);
  dp_count_dist.resize(N + 1);
  dp_count_node.resize(N + 1);
  rep(i, N - 1) {
    cin >> a >> b;
    path[a].insert(b);
    path[b].insert(a);
  }
  rrep(i, N) { cout << count_dist_subtree(i, 0) << endl; }
}


Error: TLE

Fixed submission:

#ifdef BUILT_LOCAL

#include "niklib.h"

#endif

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>

using namespace std;

typedef long long ll;
typedef unsigned long long ull;
constexpr ll mod = 1000000007;

typedef long long num;

#define bs(n) (1ll << (ll)(n))
#define rep(i, n) for (ll(i) = 0; (i) < (n); (i)++)
#define rrep(i, n) for (ll(i) = 1; (i) <= (n); (i)++)

void extgcd(num a, num b, num &ra, num &rb) {
  if (b == 0) {
    ra = a;
    rb = 0;
  } else {
    num ta, tb;
    extgcd(b, a % b, ta, tb);
    ra = tb;
    rb = ta - tb * (a / b);
  }
}

num inv_mod(num a) {
  num ta, tb;
  extgcd(a, mod, ta, tb);
  return ((ta % mod) + mod) % mod;
}

num dp_factor[300000];

num factor(num n) {
  if (n >= 300000 || n < 0)
    throw;
  if (n == 0)
    return 1;
  if (dp_factor[n] == 0)
    dp_factor[n] = (factor(n - 1) * n) % mod;
  return dp_factor[n];
}

num dp_factor_inv[300000];

num factor_inv(num n) {
  if (n >= 300000 || n < 0)
    throw;
  if (n == 0)
    return 1;
  if (dp_factor_inv[n] == 0)
    dp_factor_inv[n] = (factor_inv(n - 1) * inv_mod(n)) % mod;
  return dp_factor_inv[n];
}

num N, a, b;
vector<set<num>> path;

num count_node_subtree(num node, num parent) {
  num result = 1;
  for (auto &e : path[node]) {
    if (e == parent)
      continue;
    result += count_node_subtree(e, node);
  }
  return result;
}

vector<map<num, num>> dp_count_dist;

num count_dist_subtree(num node, num parent) {
  if (dp_count_dist[node].find(parent) == dp_count_dist[node].end()) {
    if (dp_count_dist[node].empty()) {
      num r = factor(count_node_subtree(node, parent) - 1);
      for (auto &e : path[node]) {
        if (e == parent)
          continue;
        r = (r * count_dist_subtree(e, node)) % mod;
        r = (r * factor_inv(count_node_subtree(e, node))) % mod;
      }
      dp_count_dist[node][parent] = r;
    } else {
      num r_parent = dp_count_dist[node].begin()->first;
      num r_count = dp_count_dist[node].begin()->second;
      num r = r_count;
      r = (r * (factor_inv(count_node_subtree(node, r_parent) - 1))) % mod;
      r = (r * (factor(count_node_subtree(node, parent) - 1))) % mod;
      if (r_parent != 0) {
        r = (r * factor_inv(count_node_subtree(r_parent, node))) % mod;
        r = (r * count_dist_subtree(r_parent, node)) % mod;
      }
      if (parent != 0) {
        r = (r * factor(count_node_subtree(parent, node))) % mod;
        r = (r * inv_mod(count_dist_subtree(parent, node))) % mod;
      }
      dp_count_dist[node][parent] = r;
    }
  }
  return dp_count_dist[node][parent];
}

int main() {
  cin >> N;
  path.resize(N + 1);
  dp_count_dist.resize(N + 1);
  rep(i, N - 1) {
    cin >> a >> b;
    path[a].insert(b);
    path[b].insert(a);
  }
  rrep(i, N) { cout << count_dist_subtree(i, 0) << endl; }
}