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

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, b) for (int i = (a), _ed = (b); i <= _ed; ++i)
#define DREP(i, a, b) for (int i = (a), _ed = (b); i >= _ed; --i)
#define mp(x, y) make_pair((x), (y))
#define sz(x) (int)(x).size()
#define pb push_back
typedef long long ll;
typedef pair<int, int> pii;
inline int read() {
  register int x = 0, f = 1;
  register char ch = getchar();
  while (!isdigit(ch)) {
    if (ch == '-')
      f = 0;
    ch = getchar();
  }
  while (isdigit(ch)) {
    x = x * 10 + (ch ^ '0');
    ch = getchar();
  }
  return f ? x : -x;
}

const int N = 2e5 + 5, mod = 998244353;
int n, c[2], mxd[N], f[N][3][3], g[3][3];
vector<int> E[N];
inline void inc(int &x, int y) { x = x + y < mod ? x + y : x + y - mod; }

namespace getcore {
int d[N], fa[N], stk[N], tp, tmp;
void dfs(int u, int pa, int dis, int &rt) {
  d[u] = dis, fa[u] = pa;
  if (d[u] > d[rt])
    rt = u;
  for (int v : E[u]) {
    if (v == pa)
      continue;
    dfs(v, u, dis + 1, rt);
  }
}
void Main() {
  dfs(1, 0, 0, tmp);
  dfs(tmp, 0, 0, tmp);
  while (tmp)
    stk[++tp] = tmp, tmp = fa[tmp];
  c[0] = stk[(tp + 1) >> 1];
  if (~tp & 1)
    c[1] = stk[((tp + 1) >> 1) +
               1]; // 想清楚啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊
}
} // namespace getcore

int dfs(int u, int pa, int dep) {
  mxd[u] = dep;
  for (int v : E[u]) {
    if (v == pa || v == c[0] || v == c[1])
      continue;
    mxd[u] = max(mxd[u], dfs(v, u, dep + 1));
  }
  if (mxd[u] == dep)
    f[u][1][1] = 1;
  else
    f[u][0][0] = 1;
  for (int v : E[u]) {
    if (v == pa || v == c[0] || v == c[1])
      continue;
    REP(i, 0, 2) REP(j, 0, 2) g[i][j] = 0;
    REP(a, 0, 2) REP(b, 0, 2) REP(p, 0, 2) REP(q, 0, 2) {
      int flg = mxd[u] == mxd[v];
      int s = a, t = b, val = 1ll * f[u][a][b] * f[v][p][q] % mod;
      inc(g[s][t], val);
      s = flg ? min(2, a + p) : a, t = b;
      inc(g[s][t], val);
      s = a, t = flg ? min(2, b + q) : b;
      inc(g[s][t], val);
    }
    REP(i, 0, 2) REP(j, 0, 2) f[u][i][j] = g[i][j];
  }
  return mxd[u];
}

int main() {
  // freopen("in.in","r",stdin);
  // freopen("out.out","w",stdout);
  n = read();
  REP(i, 1, n - 1) {
    int u = read(), v = read();
    E[u].pb(v), E[v].pb(u);
  }
  getcore::Main();
  int sum = 0;
  dfs(c[0], 0, 0);
  if (!c[1])
    sum = f[c[0]][1][1];
  else {
    dfs(c[1], 0, 0);
    REP(i, 0, 2)
    REP(j, 0, 2) inc(sum, 1ll * f[c[0]][1][i] * f[c[1]][j][1] % mod);
    REP(i, 0, 2)
    REP(j, 0, 2) inc(sum, 1ll * f[c[0]][i][1] * f[c[1]][1][j] % mod);
  }
  sum = 1ll * sum * (mod + 1) / 2 % mod;
  printf("%d\\n", sum);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define REP(i, a, b) for (int i = (a), _ed = (b); i <= _ed; ++i)
#define DREP(i, a, b) for (int i = (a), _ed = (b); i >= _ed; --i)
#define mp(x, y) make_pair((x), (y))
#define sz(x) (int)(x).size()
#define pb push_back
typedef long long ll;
typedef pair<int, int> pii;
inline int read() {
  register int x = 0, f = 1;
  register char ch = getchar();
  while (!isdigit(ch)) {
    if (ch == '-')
      f = 0;
    ch = getchar();
  }
  while (isdigit(ch)) {
    x = x * 10 + (ch ^ '0');
    ch = getchar();
  }
  return f ? x : -x;
}

const int N = 2e5 + 5, mod = 998244353;
int n, c[2], mxd[N], f[N][3][3], g[3][3];
vector<int> E[N];
inline void inc(int &x, int y) { x = x + y < mod ? x + y : x + y - mod; }

namespace getcore {
int d[N], fa[N], stk[N], tp, tmp;
void dfs(int u, int pa, int dis, int &rt) {
  d[u] = dis, fa[u] = pa;
  if (d[u] > d[rt])
    rt = u;
  for (int v : E[u]) {
    if (v == pa)
      continue;
    dfs(v, u, dis + 1, rt);
  }
}
void Main() {
  dfs(1, 0, 0, tmp);
  dfs(tmp, 0, 0, tmp);
  while (tmp)
    stk[++tp] = tmp, tmp = fa[tmp];
  c[0] = stk[(tp + 1) >> 1];
  if (~tp & 1)
    c[1] = stk[((tp + 1) >> 1) +
               1]; // 想清楚啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊啊
}
} // namespace getcore

int dfs(int u, int pa, int dep) {
  mxd[u] = dep;
  for (int v : E[u]) {
    if (v == pa || v == c[0] || v == c[1])
      continue;
    mxd[u] = max(mxd[u], dfs(v, u, dep + 1));
  }
  if (mxd[u] == dep)
    f[u][1][1] = 1;
  else
    f[u][0][0] = 1;
  for (int v : E[u]) {
    if (v == pa || v == c[0] || v == c[1])
      continue;
    REP(i, 0, 2) REP(j, 0, 2) g[i][j] = 0;
    REP(a, 0, 2) REP(b, 0, 2) REP(p, 0, 2) REP(q, 0, 2) {
      int flg = mxd[u] == mxd[v];
      int s = a, t = b, val = 1ll * f[u][a][b] * f[v][p][q] % mod;
      inc(g[s][t], val);
      s = flg ? min(2, a + p) : a, t = b;
      inc(g[s][t], val);
      s = a, t = flg ? min(2, b + q) : b;
      inc(g[s][t], val);
    }
    REP(i, 0, 2) REP(j, 0, 2) f[u][i][j] = g[i][j];
  }
  return mxd[u];
}

int main() {
  freopen("in.in", "r", stdin);
  freopen("out.out", "w", stdout);
  n = read();
  REP(i, 1, n - 1) {
    int u = read(), v = read();
    E[u].pb(v), E[v].pb(u);
  }
  getcore::Main();
  int sum = 0;
  dfs(c[0], 0, 0);
  if (!c[1])
    sum = f[c[0]][1][1];
  else {
    dfs(c[1], 0, 0);
    REP(i, 0, 2)
    REP(j, 0, 2) inc(sum, 1ll * f[c[0]][1][i] * f[c[1]][j][1] % mod);
    REP(i, 0, 2)
    REP(j, 0, 2) inc(sum, 1ll * f[c[0]][i][1] * f[c[1]][1][j] % mod);
  }
  sum = 1ll * sum * (mod + 1) / 2 % mod;
  printf("%d\\n", sum);
  return 0;
}