Problem p02728 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>
// #include <bits/stdc++.h>

using namespace std;

const int M = 1e9 + 7;
const int MOD = 1e9 + 7;
const int N = 4e5 + 10;
long long dp[N] = {0}, ans[N] = {0};
long long fact_ret[N] = {0};
int tree_size[N] = {0};
int visit[N] = {0};
struct edge {
  int v, next;
} edges[N];
int e_cnt = 0;
int head[N];

void add(int u, int v) {
  edges[e_cnt].v = v;
  edges[e_cnt].next = head[u];
  head[u] = e_cnt++;
}

void cal_fact() {
  fact_ret[0] = 1;
  for (int i = 1; i < N; ++i)
    fact_ret[i] = fact_ret[i - 1] * i % M;
}

int64_t extgcd(int64_t a, int64_t b, int64_t &x, int64_t &y) {
  int64_t d = a;
  if (b != 0) {
    d = extgcd(b, a % b, y, x);
    y -= (a / b) * x;
  } else {
    x = 1;
    y = 0;
  }
  return d;
}

int64_t inv_mod(int64_t a) {
  int64_t x, y;
  extgcd(a, MOD, x, y);
  return (MOD + x % MOD) % MOD;
}

int cal_size(int root) {
  visit[root] = 1;
  if (tree_size[root] != 0)
    return tree_size[root];
  tree_size[root] = 1;
  for (int e = head[root]; e != -1; e = edges[e].next) {
    if (visit[edges[e].v])
      continue;
    tree_size[root] += cal_size(edges[e].v);
  }
  return tree_size[root];
}

int cal_dp(int root) {
  // cout << "dp " << root << endl;
  visit[root] = 1;
  if (dp[root] != 0)
    return dp[root];
  dp[root] = fact_ret[tree_size[root] - 1];
  if (tree_size[root] == 1)
    return dp[root];
  for (int e = head[root]; e != -1; e = edges[e].next) {
    int v = edges[e].v;
    // cout << "v " << v << endl;
    if (visit[v])
      continue;
    dp[root] = dp[root] / fact_ret[tree_size[v]] % M;
    dp[root] = (dp[root] * cal_dp(v)) % M;
  }
  return dp[root];
}

void root_transfer(int root) {
  visit[root] = 1;
  for (int e = head[root]; e != -1; e = edges[e].next) {
    int u = root, v = edges[e].v;
    if (visit[v])
      continue;
    long long new_size_v = tree_size[u];
    long long new_size_u = tree_size[u] - tree_size[v];
    assert(("No zero", dp[v] != 0));
    long long new_dp_u = dp[u] * fact_ret[tree_size[v]] % M *
                         fact_ret[new_size_u - 1] % M * inv_mod(dp[v]) % M *
                         inv_mod(fact_ret[tree_size[u] - 1]) % M;
    long long new_dp_v = dp[v] * new_dp_u % M * fact_ret[new_size_v - 1] % M *
                         inv_mod(fact_ret[new_size_u]) % M *
                         inv_mod(fact_ret[tree_size[v] - 1]) % M;
    long long old_size_v = tree_size[v], old_size_u = tree_size[u],
              old_dp_u = dp[u], old_dp_v = dp[v];

    ans[v] = new_dp_v;
    dp[u] = new_dp_u;
    dp[v] = new_dp_v;
    tree_size[u] = new_size_u;
    tree_size[v] = new_size_v;
    // cout << root << ' ' << v << ' ' << new_dp_u << ' ' << new_dp_v << endl;
    root_transfer(v);
    dp[u] = old_dp_u;
    dp[v] = old_dp_v;
    tree_size[u] = old_size_u;
    tree_size[v] = old_size_u;
  }
}

int main() {
  memset(head, -1, sizeof(head));
  memset(dp, 0, sizeof(dp));
  memset(fact_ret, -1, sizeof(fact_ret));
  int n;
  cin >> n;
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    cin >> a >> b;
    add(a, b);
    add(b, a);
  }

  cal_fact();

  memset(visit, 0, sizeof(visit));
  memset(tree_size, 0, sizeof(tree_size));
  cal_size(1);
  memset(dp, 0, sizeof(dp));
  memset(visit, 0, sizeof(visit));
  cal_dp(1);

  ans[1] = dp[1];
  memset(visit, 0, sizeof(visit));
  root_transfer(1);

  for (int i = 1; i <= n; ++i)
    cout << ans[i] << endl;

  return 0;
}

Fixed submission:
#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>
// #include <bits/stdc++.h>

using namespace std;

const int M = 1e9 + 7;
const int MOD = 1e9 + 7;
const int N = 4e5 + 10;
long long dp[N] = {0}, ans[N] = {0};
long long fact_ret[N] = {0};
int tree_size[N] = {0};
int visit[N] = {0};
struct edge {
  int v, next;
} edges[N];
int e_cnt = 0;
int head[N];

void add(int u, int v) {
  edges[e_cnt].v = v;
  edges[e_cnt].next = head[u];
  head[u] = e_cnt++;
}

void cal_fact() {
  fact_ret[0] = 1;
  for (int i = 1; i < N; ++i)
    fact_ret[i] = fact_ret[i - 1] * i % M;
}

int64_t extgcd(int64_t a, int64_t b, int64_t &x, int64_t &y) {
  int64_t d = a;
  if (b != 0) {
    d = extgcd(b, a % b, y, x);
    y -= (a / b) * x;
  } else {
    x = 1;
    y = 0;
  }
  return d;
}

int64_t inv_mod(int64_t a) {
  int64_t x, y;
  extgcd(a, MOD, x, y);
  return (MOD + x % MOD) % MOD;
}

int cal_size(int root) {
  visit[root] = 1;
  if (tree_size[root] != 0)
    return tree_size[root];
  tree_size[root] = 1;
  for (int e = head[root]; e != -1; e = edges[e].next) {
    if (visit[edges[e].v])
      continue;
    tree_size[root] += cal_size(edges[e].v);
  }
  return tree_size[root];
}

int cal_dp(int root) {
  // cout << "dp " << root << endl;
  visit[root] = 1;
  if (dp[root] != 0)
    return dp[root];
  dp[root] = fact_ret[tree_size[root] - 1];
  if (tree_size[root] == 1)
    return dp[root];
  for (int e = head[root]; e != -1; e = edges[e].next) {
    int v = edges[e].v;
    // cout << "v " << v << endl;
    if (visit[v])
      continue;
    dp[root] = dp[root] * inv_mod(fact_ret[tree_size[v]]) % M;
    dp[root] = (dp[root] * cal_dp(v)) % M;
  }
  return dp[root];
}

void root_transfer(int root) {
  visit[root] = 1;
  for (int e = head[root]; e != -1; e = edges[e].next) {
    int u = root, v = edges[e].v;
    if (visit[v])
      continue;
    long long new_size_v = tree_size[u];
    long long new_size_u = tree_size[u] - tree_size[v];
    assert(("No zero", dp[v] != 0));
    long long new_dp_u = dp[u] * fact_ret[tree_size[v]] % M *
                         fact_ret[new_size_u - 1] % M * inv_mod(dp[v]) % M *
                         inv_mod(fact_ret[tree_size[u] - 1]) % M;
    long long new_dp_v = dp[v] * new_dp_u % M * fact_ret[new_size_v - 1] % M *
                         inv_mod(fact_ret[new_size_u]) % M *
                         inv_mod(fact_ret[tree_size[v] - 1]) % M;
    long long old_size_v = tree_size[v], old_size_u = tree_size[u],
              old_dp_u = dp[u], old_dp_v = dp[v];

    ans[v] = new_dp_v;
    dp[u] = new_dp_u;
    dp[v] = new_dp_v;
    tree_size[u] = new_size_u;
    tree_size[v] = new_size_v;
    // cout << root << ' ' << v << ' ' << new_dp_u << ' ' << new_dp_v << endl;
    root_transfer(v);
    dp[u] = old_dp_u;
    dp[v] = old_dp_v;
    tree_size[u] = old_size_u;
    tree_size[v] = old_size_u;
  }
}

int main() {
  memset(head, -1, sizeof(head));
  memset(dp, 0, sizeof(dp));
  memset(fact_ret, -1, sizeof(fact_ret));
  int n;
  cin >> n;
  for (int i = 0; i < n - 1; ++i) {
    int a, b;
    cin >> a >> b;
    add(a, b);
    add(b, a);
  }

  cal_fact();

  memset(visit, 0, sizeof(visit));
  memset(tree_size, 0, sizeof(tree_size));
  cal_size(1);
  memset(dp, 0, sizeof(dp));
  memset(visit, 0, sizeof(visit));
  cal_dp(1);

  ans[1] = dp[1];
  memset(visit, 0, sizeof(visit));
  root_transfer(1);

  for (int i = 1; i <= n; ++i)
    cout << ans[i] << endl;

  return 0;
}