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

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

#define MOD 998244353
template <typename ty1, typename ty2> inline int add(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x += y;
  return x < MOD ? x : x - MOD;
}
template <typename ty1, typename ty2> inline void addto(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x += y;
  if (x >= MOD)
    x -= MOD;
}
template <typename ty1, typename ty2> inline int sub(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x -= y;
  return x < 0 ? x + MOD : x;
}
template <typename ty1, typename ty2> inline void subto(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x -= y;
  if (x < 0)
    x += MOD;
}
template <typename ty1, typename ty2> inline int mul(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  return 1ll * x * y % MOD;
}
template <typename ty1, typename ty2> void multo(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x = 1ll * x * y % MOD;
}

long long int gcd(long long int a, long long int b) {
  if (a > b) {
    swap(a, b);
  }
  while (a) {
    swap(a, b);
    a %= b;
  }
  return b;
}
long long int lcm(long long int a, long long int b) {
  return a / gcd(a, b) * b;
}
long long int ppow(long long int i, long long int j) {
  long long int res = 1LL;
  while (j) {
    if ((j & 1LL)) {
      res *= i;
      if (res >= MOD) {
        res %= MOD;
      }
    }
    j >>= 1;
    i *= i;
    if (i >= MOD) {
      i %= MOD;
    }
  }
  return res;
}
class Combination {

public:
  vector<long long int> k;
  vector<long long int> r;
  void resize(int N) {
    k.resize(N + 2);
    r.resize(N + 2);
    k[0] = 1;
    for (int i = 1; i < N + 2; i++) {
      k[i] = k[i - 1];
      k[i] *= i;
      if (k[i] >= MOD)
        k[i] %= MOD;
    }
    long long int al = k[k.size() - 1];
    long long int iv = ppow(k[k.size() - 1], MOD - 2);
    r[k.size() - 1] = iv;
    for (int i = (int)(r.size()) - 2; i >= 0; i--) {
      r[i] = r[i + 1] * (i + 1);
      if (r[i] >= MOD) {
        r[i] %= MOD;
      }
    }
  }
  long long int C(int a, int b) {
    if (a < b)
      return 0;
    long long int up = k[a];
    long long int dw = r[b] * r[a - b];
    dw %= MOD;
    up *= dw;
    up %= MOD;
    return up;
  }
  long long int H(int a, int b) { return C(a + b - 1, b); }
  long long int catalan_number(int n) {
    return (C(2 * n, n) + MOD - C(2 * n, n - 1)) % MOD;
  }
};
Combination C;

#define MAX 200002
int n;

struct treediameter {
  // 直径は色々いい性質がある
  // 中心から(２頂点の場合それをくっつける）同じ距離のびた２つの頂点間が直径
  // 直径は最大長であることにも注意
  const vector<vector<int>> &g;
  int r1, r2;
  vector<int> d1, d2; // 両端からのdist
  int dia;
  void dfs(int v, int p, int cur, vector<int> &d) {
    d[v] = cur;
    for (auto e : g[v])
      if (e != p)
        dfs(e, v, cur + 1, d);
  }
  treediameter(const vector<vector<int>> &gg)
      : g(gg), d1(g.size()), d2(g.size()) {
    dfs(0, -1, 0, d1);
    r1 = max_element(d1.begin(), d1.end()) - d1.begin();
    dfs(r1, -1, 0, d1);
    r2 = max_element(d1.begin(), d1.end()) - d1.begin();
    dfs(r2, -1, 0, d2);
    dia = d1[r2];
  }
  vector<int> find_centre() { // 直径の中心をさがす 1頂点 or
                              // 2頂点　直径は必ずこれらすべてを通る
    vector<int> ret;
    for (int i = 0; i < d1.size(); i++) {
      if (d1[i] + d2[i] == dia && abs(d1[i] - d2[i]) <= 1) {
        ret.push_back(i);
      }
    }
    return ret;
  }
};
// treediameter(v)

vector<vector<int>> v;

bool flag[MAX];

// 0:none with only one type of operation
// 1:positive without substraction
//  2: none with mix
int dp[MAX][3];
bool us[MAX][3];
int target;
inline int dfs(int b, int ty, int pr = -1, int d = 0) {
  if (us[b][ty]) {
    return dp[b][ty];
  }
  us[b][ty] = true;
  int ava = 0;
  vector<int> can;
  for (int go : v[b]) {
    if (go == pr || flag[go])
      continue;
    ava++;
    can.push_back(go);
  }
  assert(d <= target);
  if (ava == 0) {
    // leaf
    if (ty == 1) {
      if (d == target) {
        dp[b][ty] = 1;
        return 1;
      }
      dp[b][ty] = 0;
      return 0;
    }
    if (ty == 0) {
      if (d == target) {
        dp[b][ty] = 0;
        return 0;
      }
      dp[b][ty] = 1;
      return 1;
    }
    if (ty == 2) {
      dp[b][ty] = 1;
      return 1;
    }
    exit(1);
  }
  if (ty == 1) {
    if (true) {
      for (int go : can) {
        auto ret = dfs(go, 1, b, d + 1);
        for (int goo : can) {
          if (go != goo) {
            multo(ret,
                  add(dfs(goo, 0, b, d + 1), mul(2, dfs(goo, 2, b, d + 1))));
          }
        }
        addto(dp[b][ty], ret);
      }
    } else {
      for (int go : can) {
        auto ret = dfs(go, 1, b, d + 1);
        for (int goo : can) {
          if (go != goo) {
            multo(ret, add(mul(2, dfs(goo, 0, b, d + 1)),
                           mul(1, dfs(goo, 2, b, d + 1))));
          }
        }
        addto(dp[b][ty], ret);
      }
    }
  }
  if (ty == 2) {
    dp[b][ty] = 1;
    for (int goo : can) {
      multo(dp[b][ty], mul(3, dfs(goo, 2, b, d + 1)));
    }
  }
  if (ty == 0) {
    dp[b][ty] = 1;
    for (int goo : can) {
      multo(dp[b][ty],
            add(dfs(goo, 0, b, d + 1), mul(2, dfs(goo, 2, b, d + 1))));
    }
  }
  return dp[b][ty];
}

int main() {
  cin >> n;
  v.resize(n);
  for (int i = 1; i < n; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    v[a].push_back(b);
    v[b].push_back(a);
  }
  auto dia = treediameter(v);
  auto centre = dia.find_centre();
  for (auto el : centre) {
    flag[el] = true;
  }
  if (centre.size() == 1) {
    int pos = centre[0];
    target = dia.dia / 2;
    long long int ans = 0;
    vector<int> ml;
    int overall = 1;
    for (int el3 : v[pos]) {
      ml.push_back(
          add(mul(2, dfs(el3, 0, pos, 1)), mul(1, dfs(el3, 2, pos, 1))));
      multo(overall, ml.back());
    }
    int tmp = 0;
    for (int i = 0; i < v[pos].size(); i++) {
      auto ret1 = mul(dfs(v[pos][i], 1, pos, 1), ppow(ml[i], MOD - 2));
      addto(ans, mul(tmp, ret1));
      addto(tmp, ret1);
    }
    multo(ans, overall);
    ans %= MOD;
    printf("%lld\\n", ans);
  } else {
    cerr << centre[0] << " " << centre[1] << endl;
    cerr << dia.dia << endl;
    target = dia.dia / 2;
    long long int ans = 0;
    cerr << target << " " << dfs(centre[0], 1) << " " << dfs(centre[1], 1)
         << endl;
    addto(ans, mul(dfs(centre[0], 1), dfs(centre[1], 1)));
    printf("%lld\\n", ans);
  }
  return 0;
}


Error: TLE

Fixed submission:

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

#define MOD 998244353
template <typename ty1, typename ty2> inline int add(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x += y;
  return x < MOD ? x : x - MOD;
}
template <typename ty1, typename ty2> inline void addto(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x += y;
  if (x >= MOD)
    x -= MOD;
}
template <typename ty1, typename ty2> inline int sub(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x -= y;
  return x < 0 ? x + MOD : x;
}
template <typename ty1, typename ty2> inline void subto(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x -= y;
  if (x < 0)
    x += MOD;
}
template <typename ty1, typename ty2> inline int mul(ty1 x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  return 1ll * x * y % MOD;
}
template <typename ty1, typename ty2> void multo(ty1 &x, ty2 y) {
  if (y >= MOD)
    y %= MOD;
  if (x >= MOD)
    x %= MOD;
  x = 1ll * x * y % MOD;
}

long long int gcd(long long int a, long long int b) {
  if (a > b) {
    swap(a, b);
  }
  while (a) {
    swap(a, b);
    a %= b;
  }
  return b;
}
long long int lcm(long long int a, long long int b) {
  return a / gcd(a, b) * b;
}
long long int ppow(long long int i, long long int j) {
  long long int res = 1LL;
  while (j) {
    if ((j & 1LL)) {
      res *= i;
      if (res >= MOD) {
        res %= MOD;
      }
    }
    j >>= 1;
    i *= i;
    if (i >= MOD) {
      i %= MOD;
    }
  }
  return res;
}
class Combination {

public:
  vector<long long int> k;
  vector<long long int> r;
  void resize(int N) {
    k.resize(N + 2);
    r.resize(N + 2);
    k[0] = 1;
    for (int i = 1; i < N + 2; i++) {
      k[i] = k[i - 1];
      k[i] *= i;
      if (k[i] >= MOD)
        k[i] %= MOD;
    }
    long long int al = k[k.size() - 1];
    long long int iv = ppow(k[k.size() - 1], MOD - 2);
    r[k.size() - 1] = iv;
    for (int i = (int)(r.size()) - 2; i >= 0; i--) {
      r[i] = r[i + 1] * (i + 1);
      if (r[i] >= MOD) {
        r[i] %= MOD;
      }
    }
  }
  long long int C(int a, int b) {
    if (a < b)
      return 0;
    long long int up = k[a];
    long long int dw = r[b] * r[a - b];
    dw %= MOD;
    up *= dw;
    up %= MOD;
    return up;
  }
  long long int H(int a, int b) { return C(a + b - 1, b); }
  long long int catalan_number(int n) {
    return (C(2 * n, n) + MOD - C(2 * n, n - 1)) % MOD;
  }
};
Combination C;

#define MAX 200002
int n;

struct treediameter {
  // 直径は色々いい性質がある
  // 中心から(２頂点の場合それをくっつける）同じ距離のびた２つの頂点間が直径
  // 直径は最大長であることにも注意
  const vector<vector<int>> &g;
  int r1, r2;
  vector<int> d1, d2; // 両端からのdist
  int dia;
  void dfs(int v, int p, int cur, vector<int> &d) {
    d[v] = cur;
    for (auto e : g[v])
      if (e != p)
        dfs(e, v, cur + 1, d);
  }
  treediameter(const vector<vector<int>> &gg)
      : g(gg), d1(g.size()), d2(g.size()) {
    dfs(0, -1, 0, d1);
    r1 = max_element(d1.begin(), d1.end()) - d1.begin();
    dfs(r1, -1, 0, d1);
    r2 = max_element(d1.begin(), d1.end()) - d1.begin();
    dfs(r2, -1, 0, d2);
    dia = d1[r2];
  }
  vector<int> find_centre() { // 直径の中心をさがす 1頂点 or
                              // 2頂点　直径は必ずこれらすべてを通る
    vector<int> ret;
    for (int i = 0; i < d1.size(); i++) {
      if (d1[i] + d2[i] == dia && abs(d1[i] - d2[i]) <= 1) {
        ret.push_back(i);
      }
    }
    return ret;
  }
};
// treediameter(v)

vector<vector<int>> v;

bool flag[MAX];

// 0:none with only one type of operation
// 1:positive without substraction
//  2: none with mix
int dp[MAX][3];
bool us[MAX][3];
int target;
inline int dfs(int b, int ty, int pr = -1, int d = 0) {
  if (us[b][ty]) {
    return dp[b][ty];
  }
  us[b][ty] = true;
  int ava = 0;
  vector<int> can;
  for (int go : v[b]) {
    if (go == pr || flag[go])
      continue;
    ava++;
    can.push_back(go);
  }
  assert(d <= target);
  if (ava == 0) {
    // leaf
    if (ty == 1) {
      if (d == target) {
        dp[b][ty] = 1;
        return 1;
      }
      dp[b][ty] = 0;
      return 0;
    }
    if (ty == 0) {
      if (d == target) {
        dp[b][ty] = 0;
        return 0;
      }
      dp[b][ty] = 1;
      return 1;
    }
    if (ty == 2) {
      dp[b][ty] = 1;
      return 1;
    }
    exit(1);
  }
  if (ty == 1) {
    if (true) {
      for (int go : can) {
        auto ret = dfs(go, 1, b, d + 1);
        for (int goo : can) {
          if (go != goo) {
            multo(ret,
                  add(dfs(goo, 0, b, d + 1), mul(2, dfs(goo, 2, b, d + 1))));
          }
        }
        addto(dp[b][ty], ret);
      }
    } else {
      for (int go : can) {
        auto ret = dfs(go, 1, b, d + 1);
        for (int goo : can) {
          if (go != goo) {
            multo(ret, add(mul(2, dfs(goo, 0, b, d + 1)),
                           mul(1, dfs(goo, 2, b, d + 1))));
          }
        }
        addto(dp[b][ty], ret);
      }
    }
  }
  if (ty == 2) {
    dp[b][ty] = 1;
    for (int goo : can) {
      multo(dp[b][ty], mul(3, dfs(goo, 2, b, d + 1)));
    }
  }
  if (ty == 0) {
    dp[b][ty] = 1;
    for (int goo : can) {
      multo(dp[b][ty],
            add(dfs(goo, 0, b, d + 1), mul(2, dfs(goo, 2, b, d + 1))));
    }
  }
  return dp[b][ty];
}

int main() {
  cin >> n;
  v.resize(n);
  for (int i = 1; i < n; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    v[a].push_back(b);
    v[b].push_back(a);
  }
  auto dia = treediameter(v);
  auto centre = dia.find_centre();
  for (auto el : centre) {
    flag[el] = true;
  }
  if (centre.size() == 1) {
    int pos = centre[0];
    target = dia.dia / 2;
    long long int ans = 0;
    vector<int> ml;
    int overall = 1;
    for (int el3 : v[pos]) {
      ml.push_back(
          add(mul(2, dfs(el3, 0, pos, 1)), mul(1, dfs(el3, 2, pos, 1))));
      multo(overall, ml.back());
    }
    int i = 0;
    for (int el1 : v[pos]) {
      int j = 0;
      for (int el2 : v[pos]) {
        if (el1 != el2) {
          auto ret = mul(dfs(el1, 1, pos, 1), dfs(el2, 1, pos, 1));
          multo(ret, overall);
          multo(ret, ppow(ml[i], MOD - 2));
          multo(ret, ppow(ml[j], MOD - 2));
          addto(ans, ret);
        }
        j++;
      }
      i++;
    }
    multo(ans, ppow(2, MOD - 2));
    ans %= MOD;
    printf("%lld\\n", ans);
  } else {
    cerr << centre[0] << " " << centre[1] << endl;
    cerr << dia.dia << endl;
    target = dia.dia / 2;
    long long int ans = 0;
    cerr << target << " " << dfs(centre[0], 1) << " " << dfs(centre[1], 1)
         << endl;
    addto(ans, mul(dfs(centre[0], 1), dfs(centre[1], 1)));
    printf("%lld\\n", ans);
  }
  return 0;
}