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

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

using namespace std;

#define LL long long

const int dx[4] = {1, -1, 0, 0};
const int dy[4] = {0, 0, 1, -1};

LL mod(LL x, LL md) { return ((x % md) + md) % md; }

class Mint {
public:
  static LL md;
  LL x;
  Mint(LL _x) { x = mod(_x, md); }
  Mint(const Mint &y) { x = mod(y.x, md); }
};

void MintMyExgcd(LL a, LL b, LL &x, LL &y) {
  // solve equation ax + by = 1 (mod m)
  LL md = Mint::md;
  if (b == 0) {
    y = 0;
    x = 1;
    return;
  }
  MintMyExgcd(b, a % b, x, y);
  LL t = y;
  y = mod(x - (a / b) * y, md);
  x = t;
}

LL MintNY(LL x) {
  if (x == 0) {
    return 1;
  }
  LL md = Mint::md;
  LL a, b;
  MintMyExgcd(x, md, a, b);
  while (a < 0) {
    a += md;
  }
  return a;
}

/* LL gcd(LL a, LL b) {
        if (b == 0) {
                return a;
        }
        return gcd(b, a % b);
} */

Mint operator+(Mint a, Mint b) { return Mint(mod(a.x + b.x, Mint::md)); }
Mint operator-(Mint a, Mint b) { return Mint(mod(a.x - b.x, Mint::md)); }
Mint operator*(Mint a, Mint b) { return Mint(mod(a.x * b.x, Mint::md)); }
Mint operator/(Mint a, Mint b) {
  Mint c = Mint(MintNY(b.x));
  return Mint(mod(a.x * c.x, Mint::md));
}

Mint operator-(Mint a) { return Mint(-a.x); }
bool operator==(Mint a, Mint b) { return a.x == b.x; }
bool operator!=(Mint a, Mint b) { return a.x != b.x; }
bool operator>(Mint a, Mint b) { return a.x > b.x; }
bool operator<(Mint a, Mint b) { return a.x < b.x; }
bool operator>=(Mint a, Mint b) { return a.x >= b.x; }
bool operator<=(Mint a, Mint b) { return a.x <= b.x; }

void operator+=(Mint &a, Mint b) { a = a + b; }
void operator-=(Mint &a, Mint b) { a = a - b; }
void operator*=(Mint &a, Mint b) { a = a * b; }
void operator/=(Mint &a, Mint b) { a = a / b; }

ostream &operator<<(ostream &out, Mint a) {
  out << a.x;
  return out;
}

istream &operator>>(istream &in, Mint &a) {
  in >> a.x;
  return in;
}

Mint factorial(LL n) {
  if (n == 0) {
    return Mint(1);
  }
  Mint res = 1;
  for (LL i = 1; i <= n; ++i) {
    res *= i;
  }
  return res;
}
Mint multiRange(LL s, LL t) {
  Mint res = 1;
  for (LL i = s; i <= t; ++i) {
    res *= i;
  }
  return res;
}

Mint getC(LL n, LL k) {
  if (n - k < k) {
    return getC(n, n - k);
  }
  return multiRange(n - k + 1, n) / factorial(k);
}

Mint getA(LL n, LL k) { return multiRange(n - k + 1, n); }

Mint countNonNegativeDivision(LL n, LL k) { return getC(n + k - 1, k - 1); }
Mint countPositiveDivision(LL n, LL k) {
  if (n < k) {
    return 0;
  }
  if (k <= 0) {
    return 0;
  }
  return getC(n - 1, k - 1);
}
Mint quickMod(Mint a, LL b) {
  Mint res = 1;
  while (b > 0) {
    if (b % 2) {
      res *= a;
    }
    a *= a;
    b /= 2;
  }
  return res;
}

LL Mint::md = 1000000007;

class Edge {
public:
  int t, next, idx;
  bool isF;
  Edge(int _t, int _idx, int _next = -1) : t(_t), idx(_idx), next(_next) {
    isF = false;
  }
};

vector<Mint> fac;
void makeFac(int n) {
  fac.push_back(1);
  for (int i = 1; i <= n; ++i) {
    Mint oldF = fac[i - 1];
    Mint newF = oldF * i;
    fac.push_back(newF);
  }
}
Mint DividePlan(vector<int> &a) {
  int n = 0;
  for (int i = 0; i < a.size(); ++i) {
    n += a[i];
  }
  Mint res = fac[n];
  for (int i = 0; i < a.size(); ++i) {
    res /= fac[a[i]];
  }
  return res;
}
Mint getCC(int n, int k) { return (fac[n] / fac[k]) / (fac[n - k]); }

class Tree {
public:
  int n;
  vector<Edge> e;
  vector<vector<int>> edges;
  int ec;
  int *h, *fe;
  int *s;
  vector<Mint> el, er;
  vector<int> sl, sr;
  vector<bool> elv, erv;
  bool *vis;

  void addEdge(int a, int b, int idx) {
    e.push_back(Edge(b, idx, h[a]));
    h[a] = ec++;
    e.push_back(Edge(a, idx, h[b]));
    h[b] = ec++;
  }
  vector<int> makeEdge(int a, int b) {
    vector<int> res;
    res.clear();
    res.push_back(a);
    res.push_back(b);
    return res;
  }
  Tree() {
    scanf("%d", &n);
    makeFac(n);
    h = new int[n];
    for (int i = 0; i < n; ++i) {
      h[i] = -1;
    }
    ec = 0;
    int a, b;
    for (int i = 0; i < n - 1; ++i) {
      scanf("%d%d", &a, &b);
      addEdge(a - 1, b - 1, i);
      edges.push_back(makeEdge(a - 1, b - 1));
      el.push_back(0);
      er.push_back(0);
      sl.push_back(0);
      sr.push_back(0);
      elv.push_back(false);
      erv.push_back(false);
    }
    vis = new bool[n];
    fe = new int[n];
    s = new int[n];
    for (int i = 0; i < n; ++i) {
      fe[i] = -1;
      vis[i] = false;
    }
    calcEdges(); // el, er, sl, sr
    for (int i = 0; i < n; ++i) {
      Mint res = getAns(i);
      printf("%lld\\n", res.x);
    }
  }
  void getS() {
    for (int i = 0; i < n; ++i) {
      vis[i] = false;
    }
    vis[0] = true;
    vector<int> q;
    q.push_back(0);
    int qh = 0, qt = 1;
    while (qh < qt) {
      int x = q[qh];
      for (int p = h[x]; p != -1; p = e[p].next) {
        int t = e[p].t;
        if (!vis[t]) {
          e[p].isF = false;
          q.push_back(t);
          vis[t] = true;
          ++qt;
        } else {
          e[p].isF = true;
          fe[x] = e[p].idx;
        }
      }
      ++qh;
    }
    for (int i = q.size() - 1; i >= 0; --i) {
      int x = q[i];
      s[x] = 1;
      for (int p = h[x]; p != -1; p = e[p].next) {
        if (!e[p].isF) {
          s[x] += s[e[p].t];
        }
      }
      int eIdx = fe[x];
      if (eIdx != -1) {
        if (x == edges[eIdx][0]) {
          sl[eIdx] = s[x];
          sr[eIdx] = n - s[x];
        } else {
          sr[eIdx] = s[x];
          sl[eIdx] = n - s[x];
        }
      }
    }
  }
  void getEL(int eIdx) {
    int r = edges[eIdx][0];
    // printf("getEL(eIdx = %d), r = %d\\n", eIdx, r);
    Mint res = 1;
    vector<int> div;
    for (int p = h[r]; p != -1; p = e[p].next) {
      int ni = e[p].idx;
      if (ni == eIdx) {
        continue;
      }
      bool isLeft = (r == edges[ni][0]);
      if (isLeft) {
        div.push_back(sr[ni]);
        if (!erv[ni]) {
          getER(ni);
        }
        res *= er[ni];
      } else {
        div.push_back(sl[ni]);
        if (!elv[ni]) {
          getEL(ni);
        }
        res *= el[ni];
      }
    }
    res *= DividePlan(div);
    el[eIdx] = res;
    elv[eIdx] = true;

    if (erv[eIdx]) {
      div.push_back(sr[eIdx]);
      int nn = 0;
      for (int i = 0; i < div.size(); ++i) {
        nn += div[i];
      }
      Mint total = getAns(r);
      for (int p = h[r]; p != -1; p = e[p].next) {
        int ni = e[p].idx;
        if (ni == eIdx) {
          continue;
        }
        bool isLeft = (r == edges[ni][0]);
        if (isLeft) {
          if (!elv[ni]) {
            el[ni] = (total / getCC(nn, sr[ni])) / er[ni];
            elv[ni] = true;
          }
        } else {
          if (!erv[ni]) {
            er[ni] = (total / getCC(nn, sl[ni])) / el[ni];
            erv[ni] = true;
          }
        }
      }
    }
    // cout << "res = " << res << endl;
  }
  void getER(int eIdx) {
    int r = edges[eIdx][1];
    // printf("getER(eIdx = %d), r = %d\\n", eIdx, r);
    Mint res = 1;
    vector<int> div;
    for (int p = h[r]; p != -1; p = e[p].next) {
      int ni = e[p].idx;
      if (ni == eIdx) {
        continue;
      }
      bool isLeft = (r == edges[ni][0]);
      if (isLeft) {
        div.push_back(sr[ni]);
        if (!erv[ni]) {
          getER(ni);
        }
        res *= er[ni];
      } else {
        div.push_back(sl[ni]);
        if (!elv[ni]) {
          getEL(ni);
        }
        res *= el[ni];
      }
    }
    res *= DividePlan(div);
    // cout << "res = " << res << endl;
    er[eIdx] = res;
    erv[eIdx] = true;

    if (elv[eIdx]) {
      div.push_back(sl[eIdx]);
      int nn = 0;
      for (int i = 0; i < div.size(); ++i) {
        nn += div[i];
      }
      Mint total = getAns(r);
      for (int p = h[r]; p != -1; p = e[p].next) {
        int ni = e[p].idx;
        if (ni == eIdx) {
          continue;
        }
        bool isLeft = (r == edges[ni][0]);
        if (isLeft) {
          if (!elv[ni]) {
            el[ni] = (total / getCC(nn, sr[ni])) / er[ni];
            elv[ni] = true;
          }
        } else {
          if (!erv[ni]) {
            er[ni] = (total / getCC(nn, sl[ni])) / el[ni];
            erv[ni] = true;
          }
        }
      }
    }
  }
  void getE() {
    for (int i = 0; i < n - 1; ++i) {
      if (!elv[i]) {
        getEL(i);
      }
      if (!erv[i]) {
        getER(i);
      }
    }
  }
  void calcEdges() {
    getS(); // sl, sr
    getE(); // el, er
    // el, er, sl, sr
    for (int i = 0; i < n - 1; ++i) {
      // cout << "E: " << el[i] << ' ' << er[i] << endl;
      // cout << "S: " << sl[i] << ' ' << sr[i] << endl;
    }
  }
  Mint getAns(int x) {
    vector<int> div;
    Mint res = 1;
    // div.clear();
    for (int p = h[x]; p != -1; p = e[p].next) {
      int eIdx = e[p].idx;
      bool isLeft = (edges[eIdx][0] == x);
      if (isLeft) {
        div.push_back(sr[eIdx]);
        res *= er[eIdx];
      } else {
        div.push_back(sl[eIdx]);
        res *= el[eIdx];
      }
    }
    return res * DividePlan(div);
  }
};

int main() {
  Tree t;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

using namespace std;

#define LL long long

const int dx[4] = {1, -1, 0, 0};
const int dy[4] = {0, 0, 1, -1};

LL mod(LL x, LL md) { return ((x % md) + md) % md; }

class Mint {
public:
  static LL md;
  LL x;
  Mint(LL _x) { x = mod(_x, md); }
  Mint(const Mint &y) { x = mod(y.x, md); }
};

void MintMyExgcd(LL a, LL b, LL &x, LL &y) {
  // solve equation ax + by = 1 (mod m)
  LL md = Mint::md;
  if (b == 0) {
    y = 0;
    x = 1;
    return;
  }
  MintMyExgcd(b, a % b, x, y);
  LL t = y;
  y = mod(x - (a / b) * y, md);
  x = t;
}

LL MintNY(LL x) {
  if (x == 0) {
    return 1;
  }
  LL md = Mint::md;
  LL a, b;
  MintMyExgcd(x, md, a, b);
  while (a < 0) {
    a += md;
  }
  return a;
}

/* LL gcd(LL a, LL b) {
        if (b == 0) {
                return a;
        }
        return gcd(b, a % b);
} */

Mint operator+(Mint a, Mint b) { return Mint(mod(a.x + b.x, Mint::md)); }
Mint operator-(Mint a, Mint b) { return Mint(mod(a.x - b.x, Mint::md)); }
Mint operator*(Mint a, Mint b) { return Mint(mod(a.x * b.x, Mint::md)); }
Mint operator/(Mint a, Mint b) {
  Mint c = Mint(MintNY(b.x));
  return Mint(mod(a.x * c.x, Mint::md));
}

Mint operator-(Mint a) { return Mint(-a.x); }
bool operator==(Mint a, Mint b) { return a.x == b.x; }
bool operator!=(Mint a, Mint b) { return a.x != b.x; }
bool operator>(Mint a, Mint b) { return a.x > b.x; }
bool operator<(Mint a, Mint b) { return a.x < b.x; }
bool operator>=(Mint a, Mint b) { return a.x >= b.x; }
bool operator<=(Mint a, Mint b) { return a.x <= b.x; }

void operator+=(Mint &a, Mint b) { a = a + b; }
void operator-=(Mint &a, Mint b) { a = a - b; }
void operator*=(Mint &a, Mint b) { a = a * b; }
void operator/=(Mint &a, Mint b) { a = a / b; }

ostream &operator<<(ostream &out, Mint a) {
  out << a.x;
  return out;
}

istream &operator>>(istream &in, Mint &a) {
  in >> a.x;
  return in;
}

Mint factorial(LL n) {
  if (n == 0) {
    return Mint(1);
  }
  Mint res = 1;
  for (LL i = 1; i <= n; ++i) {
    res *= i;
  }
  return res;
}
Mint multiRange(LL s, LL t) {
  Mint res = 1;
  for (LL i = s; i <= t; ++i) {
    res *= i;
  }
  return res;
}

Mint getC(LL n, LL k) {
  if (n - k < k) {
    return getC(n, n - k);
  }
  return multiRange(n - k + 1, n) / factorial(k);
}

Mint getA(LL n, LL k) { return multiRange(n - k + 1, n); }

Mint countNonNegativeDivision(LL n, LL k) { return getC(n + k - 1, k - 1); }
Mint countPositiveDivision(LL n, LL k) {
  if (n < k) {
    return 0;
  }
  if (k <= 0) {
    return 0;
  }
  return getC(n - 1, k - 1);
}
Mint quickMod(Mint a, LL b) {
  Mint res = 1;
  while (b > 0) {
    if (b % 2) {
      res *= a;
    }
    a *= a;
    b /= 2;
  }
  return res;
}

LL Mint::md = 1000000007;

class Edge {
public:
  int t, next, idx;
  bool isF;
  Edge(int _t, int _idx, int _next = -1) : t(_t), idx(_idx), next(_next) {
    isF = false;
  }
};

vector<Mint> fac;
void makeFac(int n) {
  fac.push_back(1);
  for (int i = 1; i <= n; ++i) {
    Mint oldF = fac[i - 1];
    Mint newF = oldF * i;
    fac.push_back(newF);
  }
}
Mint DividePlan(vector<int> &a) {
  int n = 0;
  for (int i = 0; i < a.size(); ++i) {
    n += a[i];
  }
  Mint res = fac[n];
  for (int i = 0; i < a.size(); ++i) {
    res /= fac[a[i]];
  }
  return res;
}

class Tree {
public:
  int n;
  vector<Edge> e;
  vector<vector<int>> edges;
  int ec;
  int *h, *fe;
  int *s;
  vector<Mint> el, er;
  vector<int> sl, sr;
  vector<bool> elv, erv;
  bool *vis;

  void addEdge(int a, int b, int idx) {
    e.push_back(Edge(b, idx, h[a]));
    h[a] = ec++;
    e.push_back(Edge(a, idx, h[b]));
    h[b] = ec++;
  }
  vector<int> makeEdge(int a, int b) {
    vector<int> res;
    res.clear();
    res.push_back(a);
    res.push_back(b);
    return res;
  }
  Tree() {
    scanf("%d", &n);
    makeFac(n);
    h = new int[n];
    for (int i = 0; i < n; ++i) {
      h[i] = -1;
    }
    ec = 0;
    int a, b;
    for (int i = 0; i < n - 1; ++i) {
      scanf("%d%d", &a, &b);
      addEdge(a - 1, b - 1, i);
      edges.push_back(makeEdge(a - 1, b - 1));
      el.push_back(0);
      er.push_back(0);
      sl.push_back(0);
      sr.push_back(0);
      elv.push_back(false);
      erv.push_back(false);
    }
    vis = new bool[n];
    fe = new int[n];
    s = new int[n];
    for (int i = 0; i < n; ++i) {
      fe[i] = -1;
      vis[i] = false;
    }
    calcEdges(); // el, er, sl, sr
    for (int i = 0; i < n; ++i) {
      Mint res = getAns(i);
      printf("%lld\\n", res.x);
    }
  }
  void getS() {
    for (int i = 0; i < n; ++i) {
      vis[i] = false;
    }
    vis[0] = true;
    vector<int> q;
    q.push_back(0);
    int qh = 0, qt = 1;
    while (qh < qt) {
      int x = q[qh];
      for (int p = h[x]; p != -1; p = e[p].next) {
        int t = e[p].t;
        if (!vis[t]) {
          e[p].isF = false;
          q.push_back(t);
          vis[t] = true;
          ++qt;
        } else {
          e[p].isF = true;
          fe[x] = e[p].idx;
        }
      }
      ++qh;
    }
    for (int i = q.size() - 1; i >= 0; --i) {
      int x = q[i];
      s[x] = 1;
      for (int p = h[x]; p != -1; p = e[p].next) {
        if (!e[p].isF) {
          s[x] += s[e[p].t];
        }
      }
      int eIdx = fe[x];
      if (eIdx != -1) {
        if (x == edges[eIdx][0]) {
          sl[eIdx] = s[x];
          sr[eIdx] = n - s[x];
        } else {
          sr[eIdx] = s[x];
          sl[eIdx] = n - s[x];
        }
      }
    }
  }
  void getEL(int eIdx) {
    int r = edges[eIdx][0];
    // printf("getEL(eIdx = %d), r = %d\\n", eIdx, r);
    Mint res = 1;
    vector<int> div;
    for (int p = h[r]; p != -1; p = e[p].next) {
      int ni = e[p].idx;
      if (ni == eIdx) {
        continue;
      }
      bool isLeft = (r == edges[ni][0]);
      if (isLeft) {
        div.push_back(sr[ni]);
        if (!erv[ni]) {
          getER(ni);
        }
        res *= er[ni];
      } else {
        div.push_back(sl[ni]);
        if (!elv[ni]) {
          getEL(ni);
        }
        res *= el[ni];
      }
    }
    res *= DividePlan(div);
    el[eIdx] = res;
    elv[eIdx] = true;
    // cout << "res = " << res << endl;
  }
  void getER(int eIdx) {
    int r = edges[eIdx][1];
    // printf("getER(eIdx = %d), r = %d\\n", eIdx, r);
    Mint res = 1;
    vector<int> div;
    for (int p = h[r]; p != -1; p = e[p].next) {
      int ni = e[p].idx;
      if (ni == eIdx) {
        continue;
      }
      bool isLeft = (r == edges[ni][0]);
      if (isLeft) {
        div.push_back(sr[ni]);
        if (!erv[ni]) {
          getER(ni);
        }
        res *= er[ni];
      } else {
        div.push_back(sl[ni]);
        if (!elv[ni]) {
          getEL(ni);
        }
        res *= el[ni];
      }
    }
    res *= DividePlan(div);
    // cout << "res = " << res << endl;
    er[eIdx] = res;
    erv[eIdx] = true;
  }
  void getE() {
    for (int i = 0; i < n - 1; ++i) {
      if (!elv[i]) {
        getEL(i);
      }
      if (!erv[i]) {
        getER(i);
      }
    }
  }
  void calcEdges() {
    getS(); // sl, sr
    getE(); // el, er
    // el, er, sl, sr
    for (int i = 0; i < n - 1; ++i) {
      // cout << "E: " << el[i] << ' ' << er[i] << endl;
      // cout << "S: " << sl[i] << ' ' << sr[i] << endl;
    }
  }
  Mint getAns(int x) {
    vector<int> div;
    Mint res = 1;
    // div.clear();
    for (int p = h[x]; p != -1; p = e[p].next) {
      int eIdx = e[p].idx;
      bool isLeft = (edges[eIdx][0] == x);
      if (isLeft) {
        div.push_back(sr[eIdx]);
        res *= er[eIdx];
      } else {
        div.push_back(sl[eIdx]);
        res *= el[eIdx];
      }
    }
    return res * DividePlan(div);
  }
};

int main() {
  Tree t;
  return 0;
}