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

#include "algorithm"
#include "iostream"
#include "map"
#include "queue"
#include "set"
#include "string"
#include "unordered_map"
#include "vector"

using namespace std;

#define ll long long
#define INF 2e9
#define MOD 1000000007

typedef vector<int> vi;
typedef pair<int, int> pii;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  bool first = true;
  for (const auto &x : v) {
    if (!first) {
      os << " ";
    }
    first = false;
    os << x;
  }
  return os;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &elem : v) {
    is >> elem;
  }
  return is;
}

ll pw(ll x, int n) {
  if (n == 1) {
    return x;
  }
  if (n % 2 == 0) {
    return pw((x * x) % MOD, n / 2);
  }
  return (pw(x, n - 1) * x) % MOD;
}

int get_min(map<int, set<int>> &m) { return m.begin()->first; }

int get_min(set<int> &m) { return *m.begin(); }

int get_max(map<int, set<int>> &m) {
  auto it = m.end();
  it--;
  return it->first;
}

int get_max(set<int> &m) {
  auto it = m.end();
  it--;
  return *it;
}

ll f[200001];
void set_f(int n) {
  f[0] = 1;
  for (int i = 1; i <= n; i++) {
    f[i] = (f[i - 1] * i) % MOD;
  }
}

vector<vi> g;
vector<ll> ans1;
vector<ll> ans2;
vector<int> cnt;

ll dv(ll a, ll b) { return (a * pw(b, MOD - 2)) % MOD; }

ll c_n_k(int n, int k) { return dv(dv(f[n], f[k]), f[n - k]); }

void dfs1(int x, int p) {
  // cout << "x=" << x << ", p=" << p << endl;
  int sum = 0;
  for (int y : g[x]) {
    if (y != p) {
      dfs1(y, x);
      sum += cnt[y];
    }
  }

  cnt[x] = sum + 1;
  ans1[x] = 1;
  if (sum == 0) {
    return;
  }

  for (int y : g[x]) {
    if (y != p) {
      ans1[x] = (ans1[x] * ans1[y]) % MOD;
      ans1[x] = (ans1[x] * c_n_k(sum, cnt[y])) % MOD;
      sum -= cnt[y];
    }
  }
}

void dfs2(int x, int p) {
  ans2[x] = 1;
  int sum = cnt[x] - 1;
  for (int y : g[x]) {
    // cout << "y=" << y << endl;
    ans2[x] = (ans2[x] * ans1[y]) % MOD;
    ans2[x] = (ans2[x] * c_n_k(sum, cnt[y])) % MOD;
    sum -= cnt[y];
  }
  ans1[x] = ans2[x];
  // cout << "x=" << x << ", p=" << p << ", ans2[x]=" << ans2[x] << endl;

  for (int y : g[x]) {
    if (y != p) {
      int cnt_x = cnt[x];
      int cnt_y = cnt[y];
      ll ans_x = ans1[x];
      ll ans_y = ans1[y];

      ans1[x] = dv(ans1[x], ans1[y]);
      ans1[x] = dv(ans1[x], c_n_k(cnt[x] - 1, cnt[y]));

      cnt[x] -= cnt[y];
      cnt[y] += cnt[x];

      dfs2(y, x);

      cnt[x] = cnt_x;
      cnt[y] = cnt_y;
      ans1[x] = ans_x;
      ans1[y] = ans_y;
    }
  }

  cnt[x] = sum + 1;
  ans1[x] = 1;
  if (sum == 0) {
    return;
  }

  for (int y : g[x]) {
    if (y != p) {
      ans1[x] = (ans1[x] * ans1[y]) % MOD;
      ans1[x] = (ans1[x] * c_n_k(sum, cnt[y])) % MOD;
      sum -= cnt[y];
    }
  }
}

void solve() {
  int n;
  cin >> n;

  set_f(n);
  g.resize(n);
  ans1.resize(n);
  ans2.resize(n);
  cnt.resize(n);

  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    g[x].push_back(y);
    g[y].push_back(x);
  }

  dfs1(0, -1);
  dfs2(0, -1);

  for (int i = 0; i < n; i++) {
    cout << ans2[i] << endl;
  }
}

int main() {
  ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);

  if (0) {
    int tests;
    cin >> tests;
    while (tests--) {
      solve();
      cout << endl;
    }
  } else {
    solve();
    cout << endl;
  }

  return 0;
}

Error: TLE

Fixed submission:

#include "algorithm"
#include "iostream"
#include "map"
#include "queue"
#include "set"
#include "string"
#include "unordered_map"
#include "vector"

using namespace std;

#define ll long long
#define INF 2e9
#define MOD 1000000007

typedef vector<int> vi;
typedef pair<int, int> pii;

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  bool first = true;
  for (const auto &x : v) {
    if (!first) {
      os << " ";
    }
    first = false;
    os << x;
  }
  return os;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (auto &elem : v) {
    is >> elem;
  }
  return is;
}

ll pw(ll x, int n) {
  if (n == 1) {
    return x;
  }
  if (n % 2 == 0) {
    return pw((x * x) % MOD, n / 2);
  }
  return (pw(x, n - 1) * x) % MOD;
}

int get_min(map<int, set<int>> &m) { return m.begin()->first; }

int get_min(set<int> &m) { return *m.begin(); }

int get_max(map<int, set<int>> &m) {
  auto it = m.end();
  it--;
  return it->first;
}

int get_max(set<int> &m) {
  auto it = m.end();
  it--;
  return *it;
}

ll f[200001];
void set_f(int n) {
  f[0] = 1;
  for (int i = 1; i <= n; i++) {
    f[i] = (f[i - 1] * i) % MOD;
  }
}

vector<vi> g;
vector<ll> ans1;
vector<ll> ans2;
vector<int> cnt;

ll dv(ll a, ll b) { return (a * pw(b, MOD - 2)) % MOD; }

ll c_n_k(int n, int k) { return dv(dv(f[n], f[k]), f[n - k]); }

void dfs1(int x, int p) {
  // cout << "x=" << x << ", p=" << p << endl;
  int sum = 0;
  for (int y : g[x]) {
    if (y != p) {
      dfs1(y, x);
      sum += cnt[y];
    }
  }

  cnt[x] = sum + 1;
  ans1[x] = 1;
  if (sum == 0) {
    return;
  }

  for (int y : g[x]) {
    if (y != p) {
      ans1[x] = (ans1[x] * ans1[y]) % MOD;
      ans1[x] = (ans1[x] * c_n_k(sum, cnt[y])) % MOD;
      sum -= cnt[y];
    }
  }
}

void dfs2(int x, int p) {
  ans2[x] = 1;
  int sum = cnt[x] - 1;
  for (int y : g[x]) {
    // cout << "y=" << y << endl;
    ans2[x] = (ans2[x] * ans1[y]) % MOD;
    ans2[x] = (ans2[x] * c_n_k(sum, cnt[y])) % MOD;
    sum -= cnt[y];
  }
  ans1[x] = ans2[x];
  // cout << "x=" << x << ", p=" << p << ", ans2[x]=" << ans2[x] << endl;

  for (int y : g[x]) {
    if (y != p) {
      int cnt_x = cnt[x];
      int cnt_y = cnt[y];
      ll ans_x = ans1[x];
      ll ans_y = ans1[y];

      cnt[x] -= cnt[y];
      cnt[y] += cnt[x];

      int sum2 = cnt[x] - 1;
      ans1[x] = 1;
      for (int yy : g[x]) {
        if (yy != y) {
          ans1[x] = (ans1[x] * ans1[yy]) % MOD;
          ans1[x] = (ans1[x] * c_n_k(sum2, cnt[yy])) % MOD;
          sum2 -= cnt[yy];
        }
      }

      dfs2(y, x);

      cnt[x] = cnt_x;
      cnt[y] = cnt_y;
      ans1[x] = ans_x;
      ans1[y] = ans_y;
    }
  }

  cnt[x] = sum + 1;
  ans1[x] = 1;
  if (sum == 0) {
    return;
  }

  for (int y : g[x]) {
    if (y != p) {
      ans1[x] = (ans1[x] * ans1[y]) % MOD;
      ans1[x] = (ans1[x] * c_n_k(sum, cnt[y])) % MOD;
      sum -= cnt[y];
    }
  }
}

void solve() {
  int n;
  cin >> n;

  set_f(n);
  g.resize(n);
  ans1.resize(n);
  ans2.resize(n);
  cnt.resize(n);

  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    g[x].push_back(y);
    g[y].push_back(x);
  }

  dfs1(0, -1);
  dfs2(0, -1);

  for (int i = 0; i < n; i++) {
    cout << ans2[i] << endl;
  }
}

int main() {
  ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);

  if (0) {
    int tests;
    cin >> tests;
    while (tests--) {
      solve();
      cout << endl;
    }
  } else {
    solve();
    cout << endl;
  }

  return 0;
}