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

// include
//------------------------------------------
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <functional>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <utility>
#include <vector>

#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>

#include <stdexcept>

using namespace std;

// type alias
using ll = long long;
using ull = unsigned long long;
using comp = complex<double>;

const ll MOD = 1000000007LL;
const double PI = 3.14159265358979323846;

// conversion
//------------------------------------------
inline int toint(string s) {
  int v;
  istringstream sin(s);
  sin >> v;
  return v;
}
template <class t> inline string tostring(t x) {
  ostringstream sout;
  sout << x;
  return sout.str();
}

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x)                                                               \\
  cerr << #x << " = " << (x) << " (l" << __line__ << ")"                       \\
       << " " << __file__ << endl

template <class Integer> class ResidueInteger {
private:
  Integer mod_;
  Integer n_;
  static Integer global_mod_;

  // computational complexity: o(log(max(a, b)))
  static inline pair<Integer, Integer> getBezoutsIdentitySolution(Integer a,
                                                                  Integer b) {
    if (b == 0) {
      return {1, 0};
    }
    auto sol = getBezoutsIdentitySolution(b, a % b);
    return {sol.second, sol.first - (a / b) * sol.second};
  }

public:
  static Integer default_mod_;

  Integer getModValue(const Integer &n) const {
    return (n % mod_ + mod_) % mod_;
  }

  // computational complexity: o(log(max(n, mod)))
  inline Integer getModInverse(const Integer &n) const {
    auto sol = getBezoutsIdentitySolution(n, mod_);
    if (n * sol.first + mod_ * sol.second != 1) {
      return -1;
    }
    return getModValue(sol.first);
  }

  ResidueInteger() : mod_(default_mod_), n_(getModValue(0)) {}
  ResidueInteger(Integer n) : mod_(default_mod_), n_(getModValue(n)) {}
  ResidueInteger(Integer n, Integer mod) : mod_(mod), n_(getModValue(n)) {}

  inline ResidueInteger operator+(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ + rhs.n_, this->mod_};
  }
  inline ResidueInteger operator+(const Integer &rhs) const {
    return {this->n_ + getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator+(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) + rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator+=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ + rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator+=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ + getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator-(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ - rhs.n_, this->mod_};
  }
  inline ResidueInteger operator-(const Integer &rhs) const {
    return {this->n_ - getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator-(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) - rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator-=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ - rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator-=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ - getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator*(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ * rhs.n_, this->mod_};
  }
  inline ResidueInteger operator*(const Integer &rhs) const {
    return {this->n_ * getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator*(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) * rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator*=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ * rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator*=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ * getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator/(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ * getModInverse(rhs.n_), this->mod_};
  }
  inline ResidueInteger operator/(const Integer &rhs) const {
    return {this->n_ * getModInverse(getModValue(rhs)), this->mod_};
  }
  friend inline ResidueInteger operator/(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) * rhs.getModInverse(rhs.n_), rhs.mod_};
  }
  inline ResidueInteger &operator/=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ * getModInverse(rhs.n_));
    return *this;
  }
  inline ResidueInteger &operator/=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ * getModInverse(getModValue(rhs)));
    return *this;
  }

  inline bool operator==(const ResidueInteger &rhs) const {
    return (this->n_ == rhs.n_) && (this->mod_ == rhs.mod_);
  }
  inline bool operator==(const Integer &rhs) const {
    return this->n_ == getModValue(rhs);
  }
  friend bool operator==(const Integer &lhs, const ResidueInteger &rhs) {
    return rhs.getModValue(lhs) == rhs.n_;
  }

  Integer get_n() const { return n_; }

  friend std::ostream &operator<<(std::ostream &lhs,
                                  const ResidueInteger &rhs) {
    return lhs << rhs.n_;
  }

  friend std::istream &operator>>(std::istream &lhs, ResidueInteger &rhs) {
    lhs >> rhs.n_;
    return lhs;
  }
};
using rll = ResidueInteger<ll>;
template <> ll rll::default_mod_ = 1000000007LL;

pair<ll, ll> calc_value(ll node, const vector<vector<ll>> &edges,
                        const vector<ll> &exclams, vector<ll> &sizes,
                        vector<ll> &mod_counts) {
  if (edges[node].size() == 0) {
    sizes[node] = 1;
    mod_counts[node] = 1;
    return {sizes[node], mod_counts[node]};
  }
  if (sizes[node] != -1) {
    return {sizes[node], mod_counts[node]};
  }

  ll node_count = 0;
  rll mod_count = 1;
  for (ll i = 0; i < edges[node].size(); ++i) {
    auto pair = calc_value(edges[node][i], edges, exclams, sizes, mod_counts);
    node_count += pair.first;
    mod_count *= pair.second;
    mod_count /= exclams[pair.first];
  }
  mod_count *= exclams[node_count];
  ++node_count;

  sizes[node] = node_count;
  mod_counts[node] = mod_count.get_n();

  return {sizes[node], mod_counts[node]};
}

void dfs(ll curr_node, const vector<vector<ll>> &edges,
         const vector<ll> &exclams, vector<ll> &sizes, vector<ll> &mod_counts,
         vector<ll> &results) {
  rll val;
  for (ll i = 0; i < edges[curr_node].size(); ++i) {
    ll next_node = edges[curr_node][i];
    ll curr_size = sizes[curr_node];
    ll curr_mod_count = sizes[curr_node];

    val = mod_counts[curr_node];
    val = val / mod_counts[next_node] / exclams[sizes.size() - 1] *
          exclams[sizes[next_node]] *
          exclams[sizes.size() - 1 - sizes[next_node]];
    mod_counts[curr_node] = val.get_n();
    sizes[curr_node] = sizes.size() - sizes[next_node];

    val = mod_counts[next_node];
    val = val * mod_counts[curr_node] * exclams[sizes.size() - 1] /
          exclams[sizes[curr_node]] /
          exclams[sizes.size() - 1 - sizes[curr_node]];
    mod_counts[next_node] = val.get_n();
    sizes[next_node] = sizes.size();

    results[next_node] = mod_counts[next_node];
    dfs(edges[curr_node][i], edges, exclams, sizes, mod_counts, results);

    sizes[next_node] = sizes.size() - sizes[curr_node];
    val = mod_counts[next_node];
    val = val / mod_counts[curr_node] / exclams[sizes.size() - 1] *
          exclams[sizes[curr_node]] *
          exclams[sizes.size() - 1 - sizes[curr_node]];
    mod_counts[next_node] = val.get_n();

    sizes[curr_node] = sizes.size();

    val = mod_counts[curr_node];
    val = val * mod_counts[next_node] * exclams[sizes.size() - 1] /
          exclams[sizes[next_node]] /
          exclams[sizes.size() - 1 - sizes[next_node]];
    mod_counts[curr_node] = val.get_n();
  }
}

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

  ll N;
  cin >> N;

  vector<ll> exclams;
  exclams.reserve(N + 1);
  rll exclam = 1;
  for (ll i = 0; i <= N; ++i) {
    exclams.emplace_back(exclam.get_n());
    exclam *= i + 1;
  }

  vector<vector<ll>> edges(N);
  ll a, b;
  for (ll i = 1; i < N; ++i) {
    cin >> a >> b;
    --a, --b;
    edges[a].emplace_back(b);
    edges[b].emplace_back(a);
  }

  vector<vector<ll>> dir_edges(N);
  {
    vector<bool> is_visited(N, false);
    queue<ll> nodes;
    nodes.push(0LL);
    is_visited[0] = true;
    while (!nodes.empty()) {
      ll curr_node = nodes.front();
      nodes.pop();
      for (auto next_node : edges[curr_node]) {
        if (is_visited[next_node])
          continue;
        dir_edges[curr_node].emplace_back(next_node);
        nodes.push(next_node);
        is_visited[next_node] = true;
      }
    }
  }

  vector<ll> sizes(N, -1);
  vector<ll> mod_counts(N);
  calc_value(0, dir_edges, exclams, sizes, mod_counts);

  vector<ll> results(N);
  {
    rll val;
    queue<ll> nodes;
    nodes.push(0LL);
    while (!nodes.empty()) {
      ll curr_node = nodes.front();
      nodes.pop();
      results[curr_node] = mod_counts[curr_node];
      for (auto next_node : dir_edges[curr_node]) {
        val = mod_counts[curr_node];
        val = val * sizes[next_node] / (N - sizes[next_node]);
        mod_counts[next_node] = val.get_n();

        sizes[next_node] = sizes.size();

        nodes.push(next_node);
      }
    }
  }

  for (ll i = 0; i < N; ++i) {
    cout << results[i] << endl;
  }

  return 0;
}

Error: TLE

Fixed submission:

// include
//------------------------------------------
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <functional>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <utility>
#include <vector>

#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <string>

#include <stdexcept>

using namespace std;

// type alias
using ll = long long;
using ull = unsigned long long;
using comp = complex<double>;

const ll MOD = 1000000007LL;
const double PI = 3.14159265358979323846;

// conversion
//------------------------------------------
inline int toint(string s) {
  int v;
  istringstream sin(s);
  sin >> v;
  return v;
}
template <class t> inline string tostring(t x) {
  ostringstream sout;
  sout << x;
  return sout.str();
}

// debug
#define DUMP(x) cerr << #x << " = " << (x) << endl
#define DEBUG(x)                                                               \\
  cerr << #x << " = " << (x) << " (l" << __line__ << ")"                       \\
       << " " << __file__ << endl

template <class Integer> class ResidueInteger {
private:
  Integer mod_;
  Integer n_;
  static Integer global_mod_;

  // computational complexity: o(log(max(a, b)))
  static inline pair<Integer, Integer> getBezoutsIdentitySolution(Integer a,
                                                                  Integer b) {
    if (b == 0) {
      return {1, 0};
    }
    auto sol = getBezoutsIdentitySolution(b, a % b);
    return {sol.second, sol.first - (a / b) * sol.second};
  }

public:
  static Integer default_mod_;

  Integer getModValue(const Integer &n) const {
    return (n % mod_ + mod_) % mod_;
  }

  // computational complexity: o(log(max(n, mod)))
  inline Integer getModInverse(const Integer &n) const {
    auto sol = getBezoutsIdentitySolution(n, mod_);
    if (n * sol.first + mod_ * sol.second != 1) {
      return -1;
    }
    return getModValue(sol.first);
  }

  ResidueInteger() : mod_(default_mod_), n_(getModValue(0)) {}
  ResidueInteger(Integer n) : mod_(default_mod_), n_(getModValue(n)) {}
  ResidueInteger(Integer n, Integer mod) : mod_(mod), n_(getModValue(n)) {}

  inline ResidueInteger operator+(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ + rhs.n_, this->mod_};
  }
  inline ResidueInteger operator+(const Integer &rhs) const {
    return {this->n_ + getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator+(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) + rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator+=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ + rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator+=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ + getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator-(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ - rhs.n_, this->mod_};
  }
  inline ResidueInteger operator-(const Integer &rhs) const {
    return {this->n_ - getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator-(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) - rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator-=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ - rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator-=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ - getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator*(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ * rhs.n_, this->mod_};
  }
  inline ResidueInteger operator*(const Integer &rhs) const {
    return {this->n_ * getModValue(rhs), this->mod_};
  }
  friend inline ResidueInteger operator*(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) * rhs.n_, rhs.mod_};
  }
  inline ResidueInteger &operator*=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ * rhs.n_);
    return *this;
  }
  inline ResidueInteger &operator*=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ * getModValue(rhs));
    return *this;
  }

  inline ResidueInteger operator/(const ResidueInteger &rhs) const {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    return {this->n_ * getModInverse(rhs.n_), this->mod_};
  }
  inline ResidueInteger operator/(const Integer &rhs) const {
    return {this->n_ * getModInverse(getModValue(rhs)), this->mod_};
  }
  friend inline ResidueInteger operator/(const Integer &lhs,
                                         const ResidueInteger &rhs) {
    return {rhs.getModValue(lhs) * rhs.getModInverse(rhs.n_), rhs.mod_};
  }
  inline ResidueInteger &operator/=(const ResidueInteger &rhs) {
    if (this->mod_ != rhs.mod_) {
      throw invalid_argument("The mods of both operands should be equal.");
    }
    this->n_ = getModValue(this->n_ * getModInverse(rhs.n_));
    return *this;
  }
  inline ResidueInteger &operator/=(const Integer &rhs) {
    this->n_ = getModValue(this->n_ * getModInverse(getModValue(rhs)));
    return *this;
  }

  inline bool operator==(const ResidueInteger &rhs) const {
    return (this->n_ == rhs.n_) && (this->mod_ == rhs.mod_);
  }
  inline bool operator==(const Integer &rhs) const {
    return this->n_ == getModValue(rhs);
  }
  friend bool operator==(const Integer &lhs, const ResidueInteger &rhs) {
    return rhs.getModValue(lhs) == rhs.n_;
  }

  Integer get_n() const { return n_; }

  friend std::ostream &operator<<(std::ostream &lhs,
                                  const ResidueInteger &rhs) {
    return lhs << rhs.n_;
  }

  friend std::istream &operator>>(std::istream &lhs, ResidueInteger &rhs) {
    lhs >> rhs.n_;
    return lhs;
  }
};
using rll = ResidueInteger<ll>;
template <> ll rll::default_mod_ = 1000000007LL;

pair<ll, ll> calc_value(ll node, const vector<vector<ll>> &edges,
                        const vector<ll> &exclams, vector<ll> &sizes,
                        vector<ll> &mod_counts) {
  if (edges[node].size() == 0) {
    sizes[node] = 1;
    mod_counts[node] = 1;
    return {sizes[node], mod_counts[node]};
  }
  if (sizes[node] != -1) {
    return {sizes[node], mod_counts[node]};
  }

  ll node_count = 0;
  rll mod_count = 1;
  for (ll i = 0; i < edges[node].size(); ++i) {
    auto pair = calc_value(edges[node][i], edges, exclams, sizes, mod_counts);
    node_count += pair.first;
    mod_count *= pair.second;
    mod_count /= exclams[pair.first];
  }
  mod_count *= exclams[node_count];
  ++node_count;

  sizes[node] = node_count;
  mod_counts[node] = mod_count.get_n();

  return {sizes[node], mod_counts[node]};
}

void dfs(ll curr_node, const vector<vector<ll>> &edges,
         const vector<ll> &exclams, vector<ll> &sizes, vector<ll> &mod_counts,
         vector<ll> &results) {
  rll val;
  for (ll i = 0; i < edges[curr_node].size(); ++i) {
    ll next_node = edges[curr_node][i];
    ll curr_size = sizes[curr_node];
    ll curr_mod_count = sizes[curr_node];

    val = mod_counts[curr_node];
    val = val / mod_counts[next_node] / exclams[sizes.size() - 1] *
          exclams[sizes[next_node]] *
          exclams[sizes.size() - 1 - sizes[next_node]];
    mod_counts[curr_node] = val.get_n();
    sizes[curr_node] = sizes.size() - sizes[next_node];

    val = mod_counts[next_node];
    val = val * mod_counts[curr_node] * exclams[sizes.size() - 1] /
          exclams[sizes[curr_node]] /
          exclams[sizes.size() - 1 - sizes[curr_node]];
    mod_counts[next_node] = val.get_n();
    sizes[next_node] = sizes.size();

    results[next_node] = mod_counts[next_node];
    dfs(edges[curr_node][i], edges, exclams, sizes, mod_counts, results);

    sizes[next_node] = sizes.size() - sizes[curr_node];
    val = mod_counts[next_node];
    val = val / mod_counts[curr_node] / exclams[sizes.size() - 1] *
          exclams[sizes[curr_node]] *
          exclams[sizes.size() - 1 - sizes[curr_node]];
    mod_counts[next_node] = val.get_n();

    sizes[curr_node] = sizes.size();

    val = mod_counts[curr_node];
    val = val * mod_counts[next_node] * exclams[sizes.size() - 1] /
          exclams[sizes[next_node]] /
          exclams[sizes.size() - 1 - sizes[next_node]];
    mod_counts[curr_node] = val.get_n();
  }
}

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

  ll N;
  cin >> N;

  vector<ll> exclams;
  exclams.reserve(N + 1);

  rll exclam = 1;
  for (ll i = 0; i <= N; ++i) {
    exclams.emplace_back(exclam.get_n());
    exclam *= i + 1;
  }

  vector<vector<ll>> edges(N);
  ll a, b;
  for (ll i = 1; i < N; ++i) {
    cin >> a >> b;
    --a, --b;
    edges[a].emplace_back(b);
    edges[b].emplace_back(a);
  }

  vector<vector<ll>> dir_edges(N);
  {
    vector<bool> is_visited(N, false);
    queue<ll> nodes;
    nodes.push(0LL);
    is_visited[0] = true;
    while (!nodes.empty()) {
      ll curr_node = nodes.front();
      nodes.pop();
      for (auto next_node : edges[curr_node]) {
        if (is_visited[next_node])
          continue;
        dir_edges[curr_node].emplace_back(next_node);
        nodes.push(next_node);
        is_visited[next_node] = true;
      }
    }
  }

  vector<ll> sizes(N, -1);
  vector<ll> mod_counts(N);
  calc_value(0, dir_edges, exclams, sizes, mod_counts);

  vector<ll> results(N);
  results[0] = mod_counts[0];
  {
    vector<bool> is_visited(N, false);
    queue<ll> nodes;
    nodes.push(0LL);
    is_visited[0] = true;
    while (!nodes.empty()) {
      ll curr_node = nodes.front();
      nodes.pop();
      for (auto next_node : edges[curr_node]) {
        if (is_visited[next_node])
          continue;
        dir_edges[curr_node].emplace_back(next_node);
        nodes.push(next_node);
        is_visited[next_node] = true;
      }
    }
  }
  dfs(0, dir_edges, exclams, sizes, mod_counts, results);

  for (ll i = 0; i < N; ++i) {
    cout << results[i] << endl;
  }

  return 0;
}