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

#include <bits/stdc++.h>
#define F(i, a, b) for (int i = a; i <= b; i++)
#define R(i, a, b) for (int i = a; i < b; i++)
#define D(i, a, b) for (int i = a; i >= b; i--)
#define go(i, x) for (int i = head[x]; i; i = Next[i])
#define pb push_back
#define si size()
#define fi first
#define se second
#define pa pair<ll, ll>
#define ll long long
#define mp make_pair
#define be begin()
#define en end()
#define put putchar('\\n')
#define gc getchar
#define pc putchar
#define N 6000055
#define mo 998244353
#define inf 1000000007
#define int ll
#define ret return puts("-1"), 0
using namespace std;
inline ll read() {
  ll su = 0, f = 1;
  char c = gc();
  for (; c < '0' || c > '9'; c = gc())
    if (c == '-')
      f = -1;
  for (; c >= '0' && c <= '9'; c = gc())
    su = su * 10 + c - '0';
  return su * f;
}
inline void wr(ll x) {
  if (x < 0)
    pc('-'), wr(-x);
  else {
    if (x >= 10)
      wr(x / 10);
    pc(x % 10 + '0');
  }
}
inline void wri(ll x) {
  wr(x);
  pc(' ');
}
inline void wrn(ll x) {
  wr(x);
  put;
}
inline void wrn(ll x, ll y) {
  wri(x);
  wrn(y);
}
inline void wrn(ll x, ll y, ll z) {
  wri(x);
  wrn(y, z);
}
int ans, n, m, f[N][3], fa[N], dis[N], lo, rt, mx;
inline void ad(int &x, int k) {
  x += k;
  x -= (x >= mo) ? mo : 0;
}
int Next[N * 2], head[N], to[N * 2], nedge, g[N][4];
#define V to[i]
void add(int a, int b) {
  Next[++nedge] = head[a];
  head[a] = nedge;
  to[nedge] = b;
}
void add_ne(int a, int b) {
  add(a, b);
  add(b, a);
}
void dfs(int x) {
  mx = (dis[x] > dis[mx]) ? x : mx;
  go(i, x) {
    if (V == fa[x])
      continue;
    fa[V] = x;
    dis[V] = dis[x] + 1;
    dfs(V);
  }
}
int ksm(ll x, int k) {
  int su = 1;
  for (; k; k >>= 1, x = x * x % mo)
    if (k & 1)
      su = su * x % mo;
  return su;
}
void solve(int x) {
  int pd = 0;
  f[x][0] = 1;
  g[x][0] = 1;
  go(i, x) {
    if (V == fa[x])
      continue;
    solve(V);
    pd++;
    g[x][3] = 1ll * g[x][3] * (((g[V][0] + g[V][1]) % mo + g[V][2]) % mo) % mo;
    ad(g[x][3], 1ll * g[x][1] * g[V][2] % mo);
    ad(g[x][3], 1ll * g[x][2] * g[V][1] % mo);
    g[x][1] = 1ll * g[x][1] * ((g[V][0] + g[V][1]) % mo) % mo;
    g[x][2] = 1ll * g[x][2] * ((g[V][0] + g[V][2]) % mo) % mo;
    ad(g[x][1], 1ll * g[x][0] * g[V][1] % mo);
    ad(g[x][2], 1ll * g[x][0] * g[V][2] % mo);
    g[x][0] = 1ll * g[x][0] * g[V][0] % mo;

    f[x][2] = 1ll * f[x][2] * ((f[V][0] + f[V][1] + f[V][2]) % mo) % mo;
    ad(f[x][2], 1ll * f[x][1] * (f[V][1] + f[V][2]) % mo);
    ad(f[x][2], 1ll * f[x][0] * f[V][2] % mo);
    f[x][1] = 1ll * f[x][1] * f[V][0] % mo;
    ad(f[x][1], 1ll * f[x][0] * f[V][1] % mo);
    f[x][0] = 1ll * f[x][0] * f[V][0] % mo;
  }
  if (pd) {
    f[x][0] = 1ll * f[x][0] * 3 % mo;
    ad(f[x][0], f[x][1] * 2 % mo);
    ad(f[x][0], f[x][2] * 2 % mo);
    g[x][0] = 1ll * g[x][0] * 3 % mo;
    ad(g[x][0], ((g[x][1] + g[x][2]) % mo * 2 % mo + g[x][3]) % mo);
    ad(g[x][1], g[x][3]);
    ad(g[x][2], g[x][3]);
  } else {
    if (dis[x] == lo / 2)
      f[x][1] = 1, f[x][0] = 2, g[x][1] = g[x][0] = g[x][2] = 1;
    else
      f[x][0] = 3, g[x][0] = 3;
  }
}
signed main() {
  n = read();
  R(i, 1, n) add_ne(read(), read());
  dfs(1);
  rt = mx;
  mx = 0;
  dis[rt] = 0;
  fa[rt] = 0;
  dfs(rt);
  lo = dis[mx];
  F(i, 1, lo / 2) mx = fa[mx];
  if (lo & 1) {
    int x = mx, y = fa[mx];
    fa[y] = x;
    dis[x] = dis[y] = 0;
    dfs(x);
    dfs(y);
    solve(x);
    solve(y);
    // wrn(f[1][0],f[1][1]);
    wrn(1ll * f[x][1] * f[y][1] % mo);
    // wrn(f[x][1],f[y][1]);
    // wrn(ksm(3,n-lo-1));
  } else {
    int x = mx, nu = 0;
    fa[x] = 0;
    dis[x] = 0;
    dfs(x);
    F(i, 1, n) if (dis[i] == lo / 2) nu++;
    // ans=(1ll*nu*(nu-1)/2)%mo*ksm(3,n-lo-1)%mo;
    solve(x);
    f[x][0] = 1;
    f[x][1] = 0;
    go(i, x) {
      // wrn(V);
      // wrn(g[V][0],g[V][1],g[V][2]);
      ans = 1ll * ans * g[V][0] % mo;
      ad(ans, 1ll * f[x][1] * f[V][1] % mo);
      f[x][1] = 1ll * f[x][1] * g[V][0] % mo;
      ad(f[x][1], 1ll * f[x][0] * f[V][1] % mo);
      f[x][0] = 1ll * f[x][0] * g[V][0] % mo;
    }
    wrn(ans);
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define F(i, a, b) for (int i = int(a); i <= int(b); i++)
#define R(i, a, b) for (int i = int(a); i < int(b); i++)
#define D(i, a, b) for (int i = int(a); i >= int(b); i--)
#define go(i, x) for (int i = head[x]; i; i = Next[i])
#define pb push_back
#define si size()
#define fi first
#define se second
#define pa pair<ll, ll>
#define ll long long
#define mp make_pair
#define be begin()
#define en end()
#define put putchar('\\n')
#define gc getchar
#define pc putchar
#define N 6000055
#define mo 998244353
#define inf 1000000007
#define ret return puts("-1"), 0
using namespace std;
inline ll read() {
  ll su = 0, f = 1;
  char c = gc();
  for (; c < '0' || c > '9'; c = gc())
    if (c == '-')
      f = -1;
  for (; c >= '0' && c <= '9'; c = gc())
    su = su * 10 + c - '0';
  return su * f;
}
inline void wr(ll x) {
  if (x < 0)
    pc('-'), wr(-x);
  else {
    if (x >= 10)
      wr(x / 10);
    pc(x % 10 + '0');
  }
}
inline void wri(ll x) {
  wr(x);
  pc(' ');
}
inline void wrn(ll x) {
  wr(x);
  put;
}
inline void wrn(ll x, ll y) {
  wri(x);
  wrn(y);
}
inline void wrn(ll x, ll y, ll z) {
  wri(x);
  wrn(y, z);
}
int ans, n, m, f[N][2], fa[N], dis[N], lo, rt, mx;
inline void ad(int &x, int k) {
  x += k;
  x -= (x >= mo) ? mo : 0;
}
int Next[N * 2], head[N], to[N * 2], nedge;
#define V to[i]
void add(int a, int b) {
  Next[++nedge] = head[a];
  head[a] = nedge;
  to[nedge] = b;
}
void add_ne(int a, int b) {
  add(a, b);
  add(b, a);
}
void dfs(int x) {
  mx = (dis[x] > dis[mx]) ? x : mx;
  go(i, x) {
    if (V == fa[x])
      continue;
    fa[V] = x;
    dis[V] = dis[x] + 1;
    dfs(V);
  }
}
int ksm(ll x, int k) {
  int su = 1;
  for (; k; k >>= 1, x = x * x % mo)
    if (k & 1)
      su = su * x % mo;
  return su;
}
void solve(int x) {
  int pd = 0;
  f[x][0] = 1;
  go(i, x) {
    if (V == fa[x])
      continue;
    solve(V);
    pd++;
    f[x][1] = 1ll * f[x][1] * f[V][0] % mo;
    ad(f[x][1], 1ll * f[x][0] * f[V][1] % mo);
    f[x][0] = 1ll * f[x][0] * f[V][0] % mo;
  }
  if (pd) {
    f[x][0] = 1ll * f[x][0] * 3 % mo;
    ad(f[x][0], f[x][1] * 2 % mo);
  } else {
    if (dis[x] == lo / 2)
      f[x][1] = 1, f[x][0] = 2;
    else
      f[x][0] = 3;
  }
}
signed main() {
  freopen("a.in", "r", stdin);
  n = read();
  R(i, 1, n) add_ne(read(), read());
  dfs(1);
  rt = mx;
  mx = 0;
  dis[rt] = 0;
  fa[rt] = 0;
  dfs(rt);
  lo = dis[mx];
  F(i, 1, lo / 2) mx = fa[mx];
  if (lo & 1) {
    int x = mx, y = fa[mx];
    fa[y] = x;
    dis[x] = dis[y] = 0;
    dfs(x);
    dfs(y);
    // solve(x);solve(y);
    wrn(ksm(3, n - lo - 1));
  } else {
    int x = mx, nu = 0;
    fa[x] = 0;
    dis[x] = 0;
    dfs(x);
    F(i, 1, n) if (dis[i] == lo / 2) nu++;
    // ans=(1ll*nu*(nu-1)/2)%mo*ksm(3,n-lo-1)%mo;
    solve(x);
    f[x][0] = 1;
    f[x][1] = 0;
    go(i, x) {
      ans = 1ll * ans * f[x][0] % mo;
      ad(ans, 1ll * f[x][1] * f[V][1] % mo);
      f[x][1] = 1ll * f[x][1] * f[V][0] % mo;
      ad(f[x][1], 1ll * f[x][0] * f[V][1] % mo);
      f[x][0] = 1ll * f[x][0] * f[V][0] % mo;
    }
    wrn(ans);
  }
  return 0;
}