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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const ll MAX_N = 200000;
const ll MOD = 1000000007;

class Prod {
public:
  vector<ll> K_;
  vector<ll> invK_;
  vector<ll> invN_;
  Prod() {
    K_.assign(MAX_N, 1);
    invK_.assign(MAX_N, 1);
    invN_.assign(MAX_N, 1);
    for (ll i = 2; i < MAX_N; i++) {
      K_.at(i) = (K_.at(i - 1) * i) % MOD;
      invN_.at(i) = MOD - ((invN_.at(MOD % i) * (MOD / i)) % MOD);
      invK_.at(i) = (invK_.at(i - 1) * invN_.at(i)) % MOD;
      //			cout << "(" << i << "," << K_.at(i) << "," <<
      //invN_.at(i) << "," << invK_.at(i) << ")" << endl;
    }
  }
  ll kaijo(ll v) { return (K_.at(v)); }
  ll inv(ll v) {
    if (v < MAX_N) {
      return (invN_.at(v));
    } else {
      return (MOD - ((inv(MOD % v) * (MOD / v)) % MOD));
    }
  }
  ll invKaijo(ll v) { return (invK_.at(v)); }
  ll comb(ll v, ll w) {
    return ((K_.at(v) * invK_.at(w) % MOD) * invK_.at(v - w) % MOD);
  }
  ll combinv(ll v, ll w) {
    return ((invK_.at(v) * K_.at(w) % MOD) * K_.at(v - w) % MOD);
  }
};

Prod p;
ll N;

class Node {
public:
  ll id_;
  pair<ll, ll> all_;
  map<Node *, pair<ll, ll>> f_; // first から入ってきたら、このノードの(個数,dp)
                                // ＝ (second.(first, second)
  void printLink() {
    for (auto it = f_.begin(); it != f_.end(); it++) {
      cout << "Link " << id_ << " to " << it->first->id_ << endl;
    }
  };
  void insertLink(Node *l) {
    f_[l].first = 0;
    f_[l].second = 0;
  }
  pair<ll, ll> setDpFirst(Node *prev) { // (house_,dp_)
    ll count = 0;
    ll prod = 1;
    for (auto it = f_.begin(); it != f_.end(); it++) {
      if (it->first != prev) {
        pair<ll, ll> next = it->first->setDpFirst(this);
        count += next.first;
        prod = (prod * next.second % MOD) * p.invKaijo(next.first) % MOD;
      }
    }
    ll mydp = (prod * p.kaijo(count)) % MOD;
    if (prev != nullptr) {
      f_[prev] = pair<ll, ll>(count + 1, mydp);
    }
    all_ = pair<ll, ll>(count, prod);
    return (pair<ll, ll>(count + 1, mydp));
  };
  pair<ll, ll> setDpSecond(Node *prev, ll prevcount,
                           ll prevdp) { // (house_,dp_) reverse
    if (prev == nullptr) {
      prevcount = 0;
      prevdp = 1;
    }
    ll count = all_.first;
    ll prod = all_.second;
    if (all_.first < N - 1) {
      for (auto it = f_.begin(); it != f_.end(); it++) {
        if (it->second.first > 0) {
          pair<ll, ll> next =
              it->first->setDpSecond(this, it->second.first, it->second.second);
          count += next.first;
          prod = (prod * next.second % MOD) * p.invKaijo(next.first) % MOD;
        }
      }
    }
    all_.first = count;
    all_.second = prod;
    ll mydp = (prod * p.kaijo(count)) % MOD;
    return (pair<ll, ll>(N - prevcount, ((all_.second * p.inv(prevdp) % MOD) *
                                         p.kaijo(prevcount) % MOD) *
                                            p.kaijo(N - prevcount - 1) % MOD));
  };
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> N;
  Node *nodes = new Node[N];
  for (int i = 0; i < N; i++) {
    nodes[i].id_ = i + 1;
    nodes[i].all_ = pair<ll, ll>(0, 0);
  }
  for (int i = 1; i <= N - 1; i++) {
    int ai, bi;
    cin >> ai >> bi;
    nodes[ai - 1].insertLink(&(nodes[bi - 1]));
    nodes[bi - 1].insertLink(&(nodes[ai - 1]));
  }
  nodes[0].setDpFirst(nullptr);
  for (ll i = 0; i < N; i++) {
    pair<ll, ll> ret = nodes[i].setDpSecond(nullptr, 0, 1);
    cout << ret.second << endl;
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const ll MAX_N = 200000;
const ll MOD = 1000000007;

class Prod {
public:
  vector<ll> K_;
  vector<ll> invK_;
  vector<ll> invN_;
  Prod() {
    K_.assign(MAX_N, 1);
    invK_.assign(MAX_N, 1);
    invN_.assign(MAX_N, 1);
    for (ll i = 2; i < MAX_N; i++) {
      K_.at(i) = (K_.at(i - 1) * i) % MOD;
      invN_.at(i) = MOD - ((invN_.at(MOD % i) * (MOD / i)) % MOD);
      invK_.at(i) = (invK_.at(i - 1) * invN_.at(i)) % MOD;
      //			cout << "(" << i << "," << K_.at(i) << "," <<
      //invN_.at(i) << "," << invK_.at(i) << ")" << endl;
    }
  }
  ll kaijo(ll v) { return (K_.at(v)); }
  ll inv(ll v) { return (invN_.at(v)); }
  ll invKaijo(ll v) { return (invK_.at(v)); }
  ll comb(ll v, ll w) {
    return ((K_.at(v) * invK_.at(w) % MOD) * invK_.at(v - w) % MOD);
  }
  ll combinv(ll v, ll w) {
    return ((invK_.at(v) * K_.at(w) % MOD) * K_.at(v - w) % MOD);
  }
};

Prod p;
ll N;

class Node {
public:
  ll id_;
  pair<ll, ll> all_;
  map<Node *, pair<ll, ll>> f_; // first から入ってきたら、このノードの(個数,dp)
                                // ＝ (second.(first, second)
  void printLink() {
    for (auto it = f_.begin(); it != f_.end(); it++) {
      cout << "Link " << id_ << " to " << it->first->id_ << endl;
    }
  };
  void insertLink(Node *l) {
    f_[l].first = 0;
    f_[l].second = 0;
  }
  pair<ll, ll> getDp(Node *prev) { // (house_,dp_)
                                   //		cout << "    counthouse " << id_ << endl;
    if (prev != nullptr) {
      if (f_[prev].second > 0) {
        return (f_[prev]);
      } // dpが計算済み
    }
    if (prev == nullptr && all_.first == N - 1) {
      //			cout << "   B " << all_.first << " " <<
      //all_.second << " " << p.kaijo(all_.first) << endl;
      return (pair<ll, ll>(all_.first + 1,
                           (all_.second * p.kaijo(all_.first)) % MOD));
    }

    // dp未計算
    ll count = 0;
    ll prod = 1;
    for (auto it = f_.begin(); it != f_.end(); it++) {
      if (it->first != prev) {
        pair<ll, ll> next = it->first->getDp(this);
        count += next.first;
        prod = (prod * next.second % MOD) * p.invKaijo(next.first) % MOD;
      }
    }
    ll mydp = (prod * p.kaijo(count)) % MOD;
    /*
                    if(prev != nullptr){
                            cout << "    " << prev->id_ << " to " << id_ << ",
       house_ = " << count + 1 << ", dp = " << mydp << endl;
                    }
                    else {
                            cout << "    " << "0" << " to " << id_ << ", house_
       = " << count + 1 << ", dp = " << mydp << endl;
                    }
    */
    if (prev != nullptr) {
      f_[prev] = pair<ll, ll>(count + 1, mydp);
    }
    all_ = pair<ll, ll>(count, prod);
    return (pair<ll, ll>(count + 1, mydp));
  };
  void setDpAll(Node *prevNode) { // (house_,dp_) reverse
    if (prevNode != nullptr) {
      pair<ll, ll> prev = prevNode->getDp(this);
      ll myHouse = all_.first + prev.first;
      ll myprod =
          (all_.second * prev.second % MOD) * p.invKaijo(prev.first) % MOD;
      all_ = pair<ll, ll>(myHouse, myprod);
      ll mydp = (all_.second * p.kaijo(all_.first)) % MOD;
      //			cout << "  A " << prevNode->id_ << " to " << id_
      //<< ", house_ = " << all_.first << ", dp = " << mydp << endl;
    }
    for (auto it = f_.begin(); it != f_.end(); it++) {
      if (it->first != prevNode) {
        it->first->setDpAll(this);
      }
    }
  };
};

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

  cin >> N;
  Node *nodes = new Node[N];
  for (int i = 0; i < N; i++) {
    nodes[i].id_ = i + 1;
    nodes[i].all_ = pair<ll, ll>(0, 0);
  }
  for (int i = 1; i <= N - 1; i++) {
    int ai, bi;
    cin >> ai >> bi;
    nodes[ai - 1].insertLink(&(nodes[bi - 1]));
    nodes[bi - 1].insertLink(&(nodes[ai - 1]));
  }

  //	for (ll i = 0; i < N; i++) {nodes[i].printLink();}cout << "end
  //printlink" << endl;

  pair<ll, ll> ret0 = nodes[0].getDp(nullptr);
  //	cout << "0: " << ret0.first << " " << ret0.second << endl;
  nodes[0].setDpAll(nullptr);
  for (ll i = 0; i < N; i++) {
    //		cout << "start " << i << endl;
    pair<ll, ll> ret = nodes[i].getDp(nullptr);
    cout << ret.second << endl;
  }
  //	cout << "END" << endl;
}