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

#include <algorithm>
#include <cmath>
#include <iostream>
#include <iterator>
#include <map>
#include <queue>
#include <unordered_map>
#include <vector>

using namespace std;

#define ll long long
#define el '\\n'
#define rep(i, a, b) for (int i = a; i < b; i++)
#define repi(i, a, b) for (int i = a; i >= b; i--)
#define MODULO 1000000007
#define VI vector<int>
#define VII vector<pair<int, int>>
#define PII pair<int, int>
#define MP make_pair
#define MAXN (int)2e5
// #define log(fmt, ...) printf(fmt, ##__VA_ARGS__)
#define log(fmt, ...)
// #define check(...) assert(##__VA_ARGS__)
#define check(...)

int extgcd(int a, int b, int &x, int &y) {
  if (a == 0) {
    x = 0;
    y = 1;
    return b;
  }
  int x1, y1;
  int d = extgcd(b % a, a, x1, y1);
  x = y1 - (b / a) * x1;
  y = x1;
  return d;
}

int invmod(int a, int m) {
  int x, y;
  int g = extgcd(a, m, x, y);
  if (g != 1) {
    return -1;
  } else {
    x = (x % m + m) % m;
    return x;
  }
}

typedef struct {
  int to;
  int edge_no;
} edge_info;

int n;
vector<edge_info> edge[MAXN];
int fac[MAXN + 1], fnv[MAXN + 1];
int ans[MAXN];

typedef struct {
  int depth;
  int parent;
  int edge_no;
} node_info;

typedef struct {
  int cnt;
  int all_nodes; // subtree node count, including root
} updown_info;

namespace my {

int mul(int a, int b) { return (ll)a * b % MODULO; }

int div(int a, int b) { return (ll)a * invmod(b, MODULO) % MODULO; }

// a * k!
int mulfac(int a, int k) { return (ll)a * fac[k] % MODULO; }

// a / k!
int divfac(int a, int k) { return (ll)a * fnv[k] % MODULO; }

} // namespace my

void solve() {
  queue<int> q;
  q.push(0);
  node_info info[n];
  rep(i, 0, n) { info[i].parent = -1; }
  info[0] = {0, 0, -1};
  while (!q.empty()) {
    int h = q.front();
    q.pop();
    for (auto x : edge[h]) {
      if (info[x.to].parent == -1) {
        q.push(x.to);
        info[x.to] = {info[h].depth + 1, h, x.edge_no};
      }
    }
  }

  int order[n];
  rep(i, 0, n) { order[i] = i; }
  auto cmp = [&info](auto i, auto j) { return info[i].depth > info[j].depth; };
  // todo: try begin(order) and end(order)
  sort(order, order + n, cmp);

  updown_info up[n];
  // up, from leaf to root
  for (auto node : order) {
    // cout << "depth(" << node << ")=" << info[node].depth << el;
    int from = info[node].parent;
    ll cnt = 1;
    int all_nodes = 1; // children count
    for (auto x : edge[node]) {
      if (x.to == from) {
        continue;
      }
      updown_info child = up[x.to];
      // cout << node << "->" << x.to << " cnt: " << child.cnt << ' '
      //      << "all nodes: " << child.all_nodes << el;

      cnt = my::mul(cnt, child.cnt);
      cnt = my::divfac(cnt, child.all_nodes);

      // cout << x.to << " fac" << fac[child.all_nodes] << el;
      // cout << x.to << " fnv" << fnv[child.all_nodes] << el;
      all_nodes += child.all_nodes;
    }
    cnt = my::mulfac(cnt, all_nodes - 1);
    up[node] = {(int)cnt, all_nodes};
    log("node=%d,cnt=%lld,all_nodes=%d\\n", node, cnt, all_nodes);
  }

  // down, from root to leaf
  reverse(order, order + n);
  updown_info down[n]; // actually the answer
  for (auto node : order) {
    int parent = info[node].parent;
    int ans1 = up[node].cnt;
    log("before ans1(%d)=%d\\n", node, ans1);

    if (parent != node) {
      int cnt_from_parent = (ll)down[parent].cnt;
      // down[parent] is the answer of parent as root
      cnt_from_parent = my::divfac(cnt_from_parent, n - 1);
      check(cnt_from_parent > 0);
      cnt_from_parent = my::mulfac(cnt_from_parent, up[node].all_nodes);
      check(cnt_from_parent > 0);
      cnt_from_parent = my::mulfac(cnt_from_parent, n - 1 - up[node].all_nodes);
      check(cnt_from_parent > 0);
      cnt_from_parent = my::div(cnt_from_parent, up[node].cnt);
      check(cnt_from_parent > 0);
      log("cnt_from_parent(%d,%d)=%d\\n", node + 1, parent + 1, cnt_from_parent);

      int all_nodes_from_parent = n - up[node].all_nodes;
      log("all_nodes_from_parent(%d,%d)=%d\\n", node + 1, parent + 1,
          all_nodes_from_parent);

      ans1 = my::divfac(ans1, up[node].all_nodes - 1);
      log("all_nodes(%d)=%d\\n", node + 1, up[node].all_nodes);
      check(ans1 > 0);
      ans1 = my::mulfac(ans1, n - 1);
      check(ans1 > 0);
      ans1 = my::mul(ans1, cnt_from_parent);
      check(ans1 > 0);
      ans1 = my::divfac(ans1, all_nodes_from_parent);
      log("ans1(%d)=%d\\n", node, ans1);
    }
    down[node] = {(int)ans1, n};
    ans[node] = ans1;
  }
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);

  cin >> n;
  rep(i, 0, n - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    edge[a].push_back({b, i});
    edge[b].push_back({a, i + n});
  }

  fac[0] = 1;
  fac[1] = 1;
  fnv[1] = 1;
  fnv[0] = 1;
  rep(i, 2, n + 1) {
    fac[i] = (ll)fac[i - 1] * i % MODULO;
    fnv[i] = invmod(fac[i], MODULO);
  }

  solve();

  rep(k, 0, n) { cout << ans[k] << el; }
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <iostream>
#include <map>
#include <queue>
#include <unordered_map>
#include <vector>

using namespace std;

#define ll long long
#define el '\\n'
#define rep(i, a, b) for (int i = a; i < b; i++)
#define repi(i, a, b) for (int i = a; i >= b; i--)
#define MODULO 1000000007
#define VI vector<int>
#define VII vector<pair<int, int>>
#define PII pair<int, int>
#define MP make_pair
#define MAXN (int)2e5
// #ifdef happy15
// #define assert(_) assert(_)
// #else
// #define assert(_)
// #endif

void bfs(int n, vector<int> edge[], int depth[], int parent[]) {
  queue<int> q;
  q.push(0);
  fill_n(depth, n, -1);
  fill_n(parent, n, -1);
  depth[0] = 0;
  parent[0] = 0;
  while (!q.empty()) {
    int h = q.front();
    q.pop();
    for (auto v : edge[h]) {
      if (parent[v] == -1) {
        depth[v] = depth[h] + 1;
        parent[v] = h;
        q.push(v);
      }
    }
  }
}
int dfs(int from, int node, int choices, vector<int> edge[], int depth[],
        int parent[]) {
  ll ans = 1;
  for (auto v : edge[node]) {
    if (v == from) {
      continue;
    }
    depth[v] = depth[node] + 1;
    parent[v] = node;
    ll ans1 =
        dfs(node, v, choices - 1 + edge[v].size() - 1, edge, depth, parent);
  }
  return ans % MODULO;
}

int dfs1(int from, int number, vector<int> edge[], vector<int> ch, bool used[],
         int n) {
  // todo: 如果ch里面没用的都是叶子，那么直接返回没有使用的个数的阶乘就行了。
  if (number > n) {
    return 1;
  }
  for (auto x : edge[from]) {
    if (!used[x]) {
      ch.push_back(x);
    }
  }
  // printf("number=%d,ch=", number);
  // for (auto x : ch) {
  //     cout << x << ',';
  // }
  // cout << el;
  ll ans = 0;
  for (auto x : ch) {
    if (!used[x]) {
      used[x] = true;
      int ans1 = dfs1(x, number + 1, edge, ch, used, n);
      ans = (ans + ans1) % MODULO;
      used[x] = false;
    }
  }
  for (auto x : edge[from]) {
    if (!used[x]) {
      ch.push_back(x);
    }
  }
  return ans;
}

int extgcd(int a, int b, int &x, int &y) {
  if (a == 0) {
    x = 0;
    y = 1;
    return b;
  }
  int x1, y1;
  int d = extgcd(b % a, a, x1, y1);
  x = y1 - (b / a) * x1;
  y = x1;
  return d;
}

int invmod(int a, int m) {
  int x, y;
  int g = extgcd(a, m, x, y);
  if (g != 1) {
    return -1;
  } else {
    x = (x % m + m) % m;
    return x;
  }
}

// unordered_map<PII, PII> dfs2_cache;
PII dfs2_cache2[2 * MAXN];

int n;
VII edge[MAXN];
int fac[MAXN + 1], fnv[MAXN + 1];

// first: ans count; second: children tree node count
PII dfs2(int from, int node, int edge_number/*, vector<PII> edge[], int n,
         int fac[], int fnv[]*/) {
  // cout << 'n' << n << el;
  if (edge_number >= 0 && dfs2_cache2[edge_number].first >= 0) {
    return dfs2_cache2[edge_number];
  }
  // PII key = MP(from, node);
  // if (dfs2_cache.find(key) != dfs2_cache.end()) {
  //     return dfs2_cache[key];
  // }
  // printf("from=%d,node=%d\\n", from, node);
  ll cnt = 1;
  int cc = 0; // children count
  for (auto x : edge[node]) {
    if (x.first == from) {
      continue;
    }
    PII ans1 = dfs2(node, x.first, x.second /*, edge, n, fac, fnv*/);
    // cout << node << "->" << x.first << ' ' << ans1.first << ' '
    //      << ans1.second << el;

    cnt = cnt * ans1.first % MODULO;
    cnt = cnt * fnv[ans1.second] % MODULO;
    // cout << "fac" << fac[ans1.second] << el;
    // cout << "fnv" << fnv[ans1.second] << el;
    cc += ans1.second;
  }
  cnt = cnt * fac[cc] % MODULO;
  // printf("from=%d,node=%d,count=%lld,children+1=%d\\n", from, node, cnt,
  //        cc + 1);
  PII ans = MP(cnt, cc + 1);
  // dfs2_cache[key] = ans;
  if (edge_number >= 0) {
    dfs2_cache2[edge_number] = ans;
  }
  return ans;
}

int solve(int k/*, int n, vector<PII> edge[], int depth[], int parent[],
          int fac[], int fnv[]*/) {
  // fill_n(depth, n, -1);
  // fill_n(parent, n, -1);
  // depth[0] = 0;
  // parent[0] = 0;
  // return dfs(0, 0, edge[0].size(), edge, depth, parent);
  // vector<int> ch;  // choices;
  // ch.push_back(k);
  // bool used[n];
  // fill_n(used, n, false);
  // used[k] = true;
  // return dfs1(k, 2, edge, ch, used, n);
  return dfs2(k, k, -1 /*, edge, n, fac, fnv*/).first;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);

  // cout << (ll) invmod(97632, MODULO) * 97632 % MODULO << el;

  // int n;
  cin >> n;
  // vector<PII> edge[n];
  rep(i, 0, n - 1) {
    int a, b;
    cin >> a >> b;
    a--, b--;
    edge[a].push_back(MP(b, i));
    edge[b].push_back(MP(a, i + n));
  }
  // int depth[n];
  // int parent[n];
  // bfs(n, edge, depth, parent);

  // int fac[n + 1], fnv[n + 1];
  fac[0] = 1;
  fac[1] = 1;
  fnv[1] = 1;
  rep(i, 2, n + 1) {
    fac[i] = (ll)fac[i - 1] * i % MODULO;
    fnv[i] = invmod(fac[i], MODULO);
    // printf("fac[%d]=%d,fnv[%d]=%d\\n", i,fac[i],i,fnv[i]);
    // assert((ll)fac[i] * fnv[i] % MODULO == 1);
  }
  fill_n(dfs2_cache2, 2 * n, MP(-1, -1));
  // rep(i,2,2*n) {
  //     cout << dfs2_cache2[i].first << ' ' << dfs2_cache2[i].second << el;
  // }
  rep(k, 0, n) {
    // cout << "---------" << el;
    cout << solve(k /*, n, edge, depth, parent, fac, fnv*/) << el;
    // printf("%d,depth=%d,parent=%d\\n", k, depth[k], parent[k]);
  }
  // rep(i, 0, 2 * n) {
  //     cout << dfs2_cache2[i].first << ' ' << dfs2_cache2[i].second << el;
  // }
}