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

#include <algorithm>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;

long long mod = 1000000007;

long long binpow(long long x, long long k) {
  if (k == 0) {
    return 1LL;
  }
  long long temp = binpow(x, k / 2);
  if (k % 2 == 1) {
    return (((temp * temp) % mod) * x) % mod;
  } else {
    return (temp * temp) % mod;
  }
}

long long modular_inverse(long long x) { return binpow(x, mod - 2); }

long long binomial(long long n, long long k) {
  long long rez = 1;
  for (long long i = n - k + 1; i <= n; i++) {
    rez *= i;
    rez %= mod;
  }
  long long temp = 1;
  for (long long i = 1; i <= k; i++) {
    temp *= i;
    temp %= mod;
  }
  rez *= modular_inverse(temp);
  rez %= mod;
  return rez;
}

int fi(int a, vector<int> &p) {
  if (p[a] != p[p[a]]) {
    p[a] = fi(p[a], p);
  }
  return p[a];
}

void uni(int a, int b, vector<int> &p) {
  int x = fi(a, p), y = fi(b, p);
  if (x != y) {
    p[x] = y;
  }
}

class SegmentTree {
public:
  class Node {
  public:
    vector<int> value;
    Node *l, *r;
    int left, right;

    Node(int x, int y) {
      value = vector<int>(26, 0);
      l = nullptr;
      r = nullptr;
      left = x;
      right = y;
    }

    void update(int p, char c) {
      if (left > p || p > right) {
        return;
      }
      if (left == p && p == right) {
        for (int i = 0; i < 26; i++) {
          value[i] = 0;
        }
        value[c - 'a'] = 1;
        return;
      }
      if (l == nullptr) {
        int mid = (left + right) / 2;
        l = new Node(left, mid);
        r = new Node(mid + 1, right);
      }
      l->update(p, c);
      r->update(p, c);
      for (int i = 0; i < 26; i++) {
        value[i] = l->value[i] + r->value[i];
      }
    }

    int query(int x, int y) {
      if (left > y || x > right) {
        return 0;
      }
      if (x <= left && right <= y) {
        int rez = 0;
        for (int i = 0; i < 26; i++) {
          if (value[i] != 0) {
            rez |= (1 << i);
          }
        }
        return rez;
      }
      if (l == nullptr) {
        return 0;
      }
      return l->query(x, y) | r->query(x, y);
    }
  };

  Node *root;
  SegmentTree(int n) { root = new Node(0, n - 1); }

  void update(int p, char c) { root->update(p, c); }

  int query(int x, int y) { return root->query(x, y); }
};

vector<long long> fact(1000001);
vector<long long> invFact(1000001);

void recalculate(int curr, int prev, vector<vector<int>> &v,
                 vector<long long> &dp, vector<int> &size) {
  long long ways = 1;
  size[curr] = 1;
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    size[curr] += size[v[curr][i]];
    ways *= dp[v[curr][i]];
    ways %= mod;
    ways *= invFact[size[v[curr][i]]];
    ways %= mod;
  }
  ways *= fact[size[curr] - 1];
  ways %= mod;
  dp[curr] = ways;
}

void calculate(int curr, int prev, vector<vector<int>> &v,
               vector<long long> &dp, vector<int> &size) {
  long long ways = 1;
  size[curr] = 1;
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    calculate(v[curr][i], curr, v, dp, size);
    size[curr] += size[v[curr][i]];
    ways *= dp[v[curr][i]];
    ways %= mod;
    ways *= invFact[size[v[curr][i]]];
    ways %= mod;
  }
  ways *= fact[size[curr] - 1];
  ways %= mod;
  dp[curr] = ways;
}

void eulerTour(int curr, int prev, vector<vector<int>> &v,
               vector<long long> &dp, vector<int> &size,
               vector<long long> &res) {
  if (prev != -1) {
    dp[prev] *= modular_inverse(dp[curr]);
    dp[prev] %= mod;
    dp[prev] *= fact[size[curr]];
    dp[prev] %= mod;
    dp[prev] *= invFact[size[prev] - 1];
    dp[prev] %= mod;
    dp[curr] *= invFact[size[curr] - 1];
    dp[curr] %= mod;
    size[prev] -= size[curr];
    size[curr] += size[prev];
    dp[prev] *= fact[size[prev] - 1];
    dp[prev] %= mod;
    dp[curr] *= fact[size[curr] - 1];
    dp[curr] %= mod;
    dp[curr] *= dp[prev];
    dp[curr] %= mod;
    dp[curr] *= invFact[size[prev]];
    dp[curr] %= mod;
  }
  res[curr] = dp[curr];
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    eulerTour(v[curr][i], curr, v, dp, size, res);
  }
  if (prev != -1) {
    dp[curr] *= modular_inverse(dp[prev]);
    dp[curr] %= mod;
    dp[curr] *= fact[size[prev]];
    dp[curr] %= mod;
    dp[curr] *= invFact[size[curr] - 1];
    dp[curr] %= mod;
    dp[prev] *= invFact[size[prev] - 1];
    dp[prev] %= mod;
    size[curr] -= size[prev];
    size[prev] += size[curr];
    dp[curr] *= fact[size[curr] - 1];
    dp[curr] %= mod;
    dp[prev] *= fact[size[prev] - 1];
    dp[prev] %= mod;
    dp[prev] *= dp[curr];
    dp[prev] %= mod;
    dp[prev] *= invFact[size[curr]];
    dp[prev] %= mod;
  }
}

int main() {
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  vector<vector<int>> v(n);
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    v[x].push_back(y);
    v[y].push_back(x);
  }
  vector<long long> dp(n);
  vector<long long> res(n);
  vector<int> size(n);
  fact[0] = 1;
  invFact[0] = 1;
  for (int i = 1; i <= n; i++) {
    fact[i] = (i * fact[i - 1]) % mod;
    invFact[i] = modular_inverse(fact[i]);
  }
  calculate(0, -1, v, dp, size);
  eulerTour(0, -1, v, dp, size, res);
  for (int i = 0; i < n; i++) {
    cout << res[i] << '\\n';
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iomanip>
#include <iostream>
#include <map>
#include <math.h>
#include <queue>
#include <set>
#include <stack>
#include <vector>

using namespace std;

long long mod = 1000000007;

long long binpow(long long x, long long k) {
  if (k == 0) {
    return 1LL;
  }
  long long temp = binpow(x, k / 2);
  if (k % 2 == 1) {
    return (((temp * temp) % mod) * x) % mod;
  } else {
    return (temp * temp) % mod;
  }
}

long long modular_inverse(long long x) { return binpow(x, mod - 2); }

long long binomial(long long n, long long k) {
  long long rez = 1;
  for (long long i = n - k + 1; i <= n; i++) {
    rez *= i;
    rez %= mod;
  }
  long long temp = 1;
  for (long long i = 1; i <= k; i++) {
    temp *= i;
    temp %= mod;
  }
  rez *= modular_inverse(temp);
  rez %= mod;
  return rez;
}

int fi(int a, vector<int> &p) {
  if (p[a] != p[p[a]]) {
    p[a] = fi(p[a], p);
  }
  return p[a];
}

void uni(int a, int b, vector<int> &p) {
  int x = fi(a, p), y = fi(b, p);
  if (x != y) {
    p[x] = y;
  }
}

class SegmentTree {
public:
  class Node {
  public:
    vector<int> value;
    Node *l, *r;
    int left, right;

    Node(int x, int y) {
      value = vector<int>(26, 0);
      l = nullptr;
      r = nullptr;
      left = x;
      right = y;
    }

    void update(int p, char c) {
      if (left > p || p > right) {
        return;
      }
      if (left == p && p == right) {
        for (int i = 0; i < 26; i++) {
          value[i] = 0;
        }
        value[c - 'a'] = 1;
        return;
      }
      if (l == nullptr) {
        int mid = (left + right) / 2;
        l = new Node(left, mid);
        r = new Node(mid + 1, right);
      }
      l->update(p, c);
      r->update(p, c);
      for (int i = 0; i < 26; i++) {
        value[i] = l->value[i] + r->value[i];
      }
    }

    int query(int x, int y) {
      if (left > y || x > right) {
        return 0;
      }
      if (x <= left && right <= y) {
        int rez = 0;
        for (int i = 0; i < 26; i++) {
          if (value[i] != 0) {
            rez |= (1 << i);
          }
        }
        return rez;
      }
      if (l == nullptr) {
        return 0;
      }
      return l->query(x, y) | r->query(x, y);
    }
  };

  Node *root;
  SegmentTree(int n) { root = new Node(0, n - 1); }

  void update(int p, char c) { root->update(p, c); }

  int query(int x, int y) { return root->query(x, y); }
};

vector<long long> fact;
vector<long long> invFact;

void recalculate(int curr, int prev, vector<vector<int>> &v,
                 vector<long long> &dp, vector<int> &size) {
  long long ways = 1;
  size[curr] = 1;
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    size[curr] += size[v[curr][i]];
    ways *= dp[v[curr][i]];
    ways %= mod;
    ways *= invFact[size[v[curr][i]]];
    ways %= mod;
  }
  ways *= fact[size[curr] - 1];
  ways %= mod;
  dp[curr] = ways;
}

void calculate(int curr, int prev, vector<vector<int>> &v,
               vector<long long> &dp, vector<int> &size) {
  long long ways = 1;
  size[curr] = 1;
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    calculate(v[curr][i], curr, v, dp, size);
    size[curr] += size[v[curr][i]];
    ways *= dp[v[curr][i]];
    ways %= mod;
    ways *= invFact[size[v[curr][i]]];
    ways %= mod;
  }
  ways *= fact[size[curr] - 1];
  ways %= mod;
  dp[curr] = ways;
}

void eulerTour(int curr, int prev, vector<vector<int>> &v,
               vector<long long> &dp, vector<int> &size,
               vector<long long> &res) {
  if (prev != -1) {
    recalculate(prev, curr, v, dp, size);
    recalculate(curr, -1, v, dp, size);
  }
  res[curr] = dp[curr];
  for (int i = 0; i < v[curr].size(); i++) {
    if (v[curr][i] == prev) {
      continue;
    }
    eulerTour(v[curr][i], curr, v, dp, size, res);
    recalculate(v[curr][i], curr, v, dp, size);
    recalculate(curr, prev, v, dp, size);
  }
}

int main() {
  ios::sync_with_stdio(false);
  int n;
  cin >> n;
  vector<vector<int>> v(n);
  for (int i = 0; i < n - 1; i++) {
    int x, y;
    cin >> x >> y;
    x--;
    y--;
    v[x].push_back(y);
    v[y].push_back(x);
  }
  vector<long long> dp(n);
  vector<long long> res(n);
  vector<int> size(n);
  fact.push_back(1);
  fact.push_back(1);
  invFact.push_back(1);
  invFact.push_back(1);
  for (int i = 2; i <= n; i++) {
    fact.push_back((i * fact[i - 1]) % mod);
    invFact.push_back(modular_inverse(fact[i]));
  }
  calculate(0, -1, v, dp, size);
  eulerTour(0, -1, v, dp, size, res);
  for (int i = 0; i < n; i++) {
    cout << res[i] << '\\n';
  }
  return 0;
}