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


#include <iostream>
#include <vector>

using namespace std;

typedef long long ll;

// auto mod int
// https://youtu.be/L8grWxBlIZ4?t=9858
// https://youtu.be/ERZuLAxZffQ?t=4807 : optimize
// https://youtu.be/8uowVvQ_-Mo?t=1329 : division
const int mod = 1000000007;
// const int mod = 998244353;
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};

struct combination {
  vector<mint> fact, ifact;
  combination(int n) : fact(n + 1), ifact(n + 1) {
    fact[0] = 1;
    for (int i = 1; i <= n; ++i)
      fact[i] = fact[i - 1] * i;
    ifact[n] = fact[n].inv();
    for (int i = n; i >= 1; --i)
      ifact[i - 1] = ifact[i] * i;
  }
  mint operator()(int n, int k) {
    if (k < 0 || k > n)
      return 0;
    return fact[n] * ifact[k] * ifact[n - k];
  }
} comb(200000);

class DP {
public:
  mint dp; // 辺の元のノードの DP値
  int t;   // 辺の元のノードより先のノード数
  DP(mint _dp = 1, int _t = 0) {
    dp = _dp;
    t = _t;
  }
};

DP vtx_dp[200000];
vector<int> v_1ring[200000];

static void CalcDp_Dfs(int parent_vid, int node_vid) {
  DP &piv_dp = vtx_dp[node_vid];

  for (auto v : v_1ring[node_vid]) {
    if (v == parent_vid)
      continue;

    CalcDp_Dfs(node_vid, v);

    DP &ch = vtx_dp[v];
    piv_dp.dp *= ch.dp;
    piv_dp.dp *= comb(piv_dp.t + ch.t, ch.t);
    piv_dp.t += ch.t;
  }

  piv_dp.t += 1;
}

// 上から下へ
static void CalcDp_Bfs(const int parent_id, const int node_id) {

  for (auto v : v_1ring[node_id]) {
    if (v == parent_id)
      continue;

    DP tmp = vtx_dp[node_id];
    // 辺 node_id - v が無かった状態に戻す
    tmp.t -= vtx_dp[v].t;
    tmp.dp /= comb(tmp.t + vtx_dp[v].t - 1, tmp.t - 1);
    tmp.dp /= vtx_dp[v].dp;

    // dfsと同様に追加する
    vtx_dp[v].dp *= tmp.dp;
    vtx_dp[v].dp *= comb(tmp.t + vtx_dp[v].t - 1, tmp.t);
    vtx_dp[v].t += tmp.t;

    CalcDp_Bfs(node_id, v);
  }
}

int main(void) {

  int n;
  cin >> n;

  vector<pair<int, int>> edges(n);

  for (int i = 0; i < n - 1; ++i) {
    cin >> edges[i].first >> edges[i].second;
    --edges[i].first;
    --edges[i].second;
  }

  for (int i = 0; i < n - 1; ++i) {
    v_1ring[edges[i].first].push_back(edges[i].second);
    v_1ring[edges[i].second].push_back(edges[i].first);
  }
  CalcDp_Dfs(-1, 0);
  CalcDp_Bfs(-1, 0);

  for (int i = 0; i < n; ++i)
    cout << vtx_dp[i].dp.x << "\\n";

  return 0;
}


Error: TLE

Fixed submission:


#include <iostream>
#include <vector>

using namespace std;

typedef long long ll;

// auto mod int
// https://youtu.be/L8grWxBlIZ4?t=9858
// https://youtu.be/ERZuLAxZffQ?t=4807 : optimize
// https://youtu.be/8uowVvQ_-Mo?t=1329 : division
const int mod = 1000000007;
// const int mod = 998244353;
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};

// tree
//  t[i] = (vid, num_ch)
//  t[i]は，頂点iに入ってくる辺 (i <-- vid)について, vidより先の子供の数を表す

static mint factrial(int n) {
  mint c = 1;
  for (int i = 1; i <= n; ++i) {
    c *= i;
  }
  return c;
}

static mint CalcValue(int parent_id, int node_id,
                      const vector<vector<pair<int, int>>> &graph_t) {
  // cout << parent_id << " " << node_id << endl;
  // for ( const auto v : graph_t[node_id] )
  //     cout << "---" << node_id << "-" << v.first << " " <<v.second << "\\n";

  mint res = 1;
  int sum = 0;
  for (const auto v : graph_t[node_id]) {
    if (v.first == parent_id)
      continue;
    sum += v.second;
    res *= CalcValue(node_id, v.first, graph_t);
  }

  res *= factrial(sum);

  for (auto v : graph_t[node_id]) {
    if (v.first == parent_id)
      continue;
    res /= factrial(v.second);
  }

  return res;
}

static int CalcT_Dfs(int parent_vid, int node_vid,
                     const vector<vector<int>> &v_1ring,
                     vector<vector<pair<int, int>>> &t) {
  int sum = 0;
  for (auto v : v_1ring[node_vid]) {
    if (v == parent_vid)
      continue;
    int num_ch = CalcT_Dfs(node_vid, v, v_1ring, t);
    t[node_vid].push_back(make_pair(v, num_ch));
    sum += num_ch;
  }
  return sum + 1;
}

static void CalcT_Dfs_Second(const int node_vid,
                             vector<vector<pair<int, int>>> &t) {
  int sum = (int)t.size();

  for (auto v : t[node_vid]) {
    CalcT_Dfs_Second(v.first, t);
  }

  for (auto v : t[node_vid]) {
    t[v.first].push_back(make_pair(node_vid, sum - v.second));
  }
}

class MyGraph {
  int m_n; // num verts  edges
  vector<pair<int, int>> m_edges;

  vector<vector<int>> m_v_1ring;
  vector<vector<pair<int, int>>> m_t;

public:
  MyGraph(int n, vector<pair<int, int>> edges) {
    m_n = n;
    m_edges = edges;
    m_v_1ring.resize(n);

    for (int i = 0; i < n - 1; ++i) {
      m_v_1ring[edges[i].first].push_back(edges[i].second);
      m_v_1ring[edges[i].second].push_back(edges[i].first);
    }

    CalcGraphT();

    for (int i = 0; i < m_n; ++i)
      cout << CalcValue(-1, i, m_t).x << "\\n";
  }

  void CalcGraphT() {
    // calc t[i]
    m_t.clear();
    m_t.resize(m_n);
    CalcT_Dfs(-1, 0, m_v_1ring, m_t);
    CalcT_Dfs_Second(0, m_t);

    // for( int i = 0; i < m_n; ++i) {
    //   for ( auto it : m_t[i]) cout << it.first << " " << it.second << " ";
    //   cout << "-----" << i << "\\n";
    // }
  }
};

int main(void) {

  int n;
  cin >> n;

  vector<pair<int, int>> edges(n);

  for (int i = 0; i < n - 1; ++i) {
    cin >> edges[i].first >> edges[i].second;
    --edges[i].first;
    --edges[i].second;
  }

  MyGraph g(n, edges);
  return 0;
}