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

#include <bits/stdc++.h>

#define ll long long
#define fornum(A, B, C) for (A = B; A < C; ++A)
#define pii pair<int, int>
#define pll pair<ll, ll>

using namespace std;

/////////////////////////////////////////////////////
#define MOD 998244353ll
ll moddiv(ll a, ll b) {
  a %= MOD;
  b %= MOD;
  for (ll m = MOD - 2; m > 0; m /= 2) {
    if (m & 1) {
      a = a * b % MOD;
    }
    b = b * b % MOD;
  }
  return a;
}

ll N, Len;
vector<ll> abv[202020];
vector<ll> abl[202020];
ll mk[202020], p3[202020];
ll dp[202020][4];
ll i, j, k;
ll ans;

void dfs(ll a, ll p, ll d) {
  if (abv[a].size() == 1) {
    dp[a][2] = 3;
    if (d == 0) {
      dp[a][0] = 2;
      dp[a][1] = 1;
    } else {
      dp[a][0] = 3;
    }
    return;
  }
  ll i;
  ll c = 1, e = 0;
  dp[a][2] = 1;
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p) {
      continue;
    }
    dfs(b, a, d - 1);
    c = c * dp[b][0] % MOD;
    dp[a][2] *= dp[b][2];
    dp[a][2] %= MOD;
  }
  dp[a][0] = (dp[a][2] * 2 + c) % MOD;
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p) {
      continue;
    }
    dp[a][1] += moddiv(c * dp[b][1], dp[b][0]);
    dp[a][1] %= MOD;
  }
  dp[a][2] *= 3;
  dp[a][2] %= MOD;
}

ll lencalc(ll a, ll p, ll d) {
  // printf("%lld,%lld,%lld\\n", p, a, d);
  ll i = 0, ret = 0;
  pll c = {0, 0};
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p)
      continue;
    if (c.first <= abl[a][i]) {
      c.second = c.first;
      c.first = abl[a][i];
    } else if (c.second <= abl[a][i]) {
      c.second = abl[a][i];
    }
  }
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p) {
      abl[a][i] = d;
      continue;
    }
    ll r = 0;
    if (c.first == abl[a][i]) {
      r = lencalc(b, a, max(d, c.second) + 1);
    } else {
      r = lencalc(b, a, max(d, c.first) + 1);
    }
    abl[a][i] = r;
    ret = max(ret, r);
  }
  return ret + 1;
}

int main() {
  scanf("%lld", &N);
  fornum(i, 1, N) {
    ll a, b;
    scanf("%lld%lld", &a, &b);
    abv[a].push_back(b);
    abl[a].push_back(0);
    abv[b].push_back(a);
    abl[b].push_back(0);
  }
  lencalc(1, 0, 0);
  lencalc(1, 0, 0);
  fornum(i, 1, N + 1) {
    if (abv[i].size() == 1) {
      Len = max(Len, abl[i][0]);
    }
  }

  if (Len % 2) {
    pll a = {0, 0};
    fornum(i, 1, N + 1) {
      fornum(j, 0, abv[i].size()) {
        if ((Len + 1) / 2 == abl[i][j]) {
          break;
        }
      }
      if (j < abv[i].size()) {
        a.first = i;
        a.second = abv[i][j];
      }
    }
    // printf("b%lld,%lld\\n", a.first, a.second);
    dfs(a.first, a.second, Len / 2);
    dfs(a.second, a.first, Len / 2);
    ans = dp[a.first][1] * dp[a.second][1] % MOD;
  } else {
    fornum(i, 1, N + 1) {
      if (abv[i].size() == 1)
        continue;
      sort(abl[i].begin(), abl[i].end(), greater<ll>());
      if (abl[i][0] == abl[i][1])
        break;
    }
    ll a = i;
    // printf("a%lld\\n", a);
    dfs(a, 0, Len / 2);
    ll c = 1;
    fornum(i, 0, abv[a].size()) {
      ll b = abv[a][i];
      // printf("cc%lld,%lld,%lld\\n", b, dp[b][0], dp[b][1]);
      c = c * (dp[b][0] * 2 - dp[b][2] + MOD) % MOD;
    }
    ll d = 0;
    fornum(i, 0, abv[a].size()) {
      ll b = abv[a][i];
      ans += moddiv(d * dp[b][1], dp[b][0] * 2 - dp[b][2] + MOD);
      ans %= MOD;
      d += moddiv(c * dp[b][1], dp[b][0] * 2 - dp[b][2] + MOD);
      d %= MOD;
    }
  }

  printf("%lld", ans);

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

#define ll long long
#define fornum(A, B, C) for (A = B; A < C; ++A)
#define pii pair<int, int>
#define pll pair<ll, ll>

using namespace std;

/////////////////////////////////////////////////////
#define MOD 998244353ll

ll N, gr;
vector<ll> abv[202020];
ll mk[202020], p3[202020];
ll i, j, k;
ll ans;

pll grcalc(ll a, ll p) {
  ll i = 0;
  pll ret = {0, 0};
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p)
      continue;
    pll p = grcalc(b, a);
    if (ret.first < p.first) {
      ret.first = p.first;
    }
    if (ret.first < p.second + ret.second) {
      ret.first = p.second + ret.second;
    }
    if (ret.second < p.second) {
      ret.second = p.second;
    }
  }
  ++ret.second;
  return ret;
}

void dfs2(ll a, ll p, ll d, ll dd, ll &ra, ll &rb) {
  ra = 1;
  rb = 1;
  ll i, j;
  if (abv[a].size() == 1) {
    if (d == dd || d == gr - dd) {
      rb = 0;
    }
    // printf("%lld:%lld,%lld\\n", a, ra, rb);
    return;
  }

  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p)
      continue;
    ll rra, rrb;
    dfs2(b, a, d + 1, dd, rra, rrb);
    ra *= rra * 3;
    ra %= MOD;
    rb *= rrb + rra * 2;
    rb %= MOD;
  }
  // printf("%lld:%lld,%lld\\n", a, ra, rb);
  return;
}
ll dfs1(ll a, ll p, ll d) {
  mk[d] = a;
  ll i, j;
  ll ret = 0;
  if (d == gr) {
    if (mk[0] > mk[d])
      return 0;
    i = 0;
    ret = 1;
    fornum(i, 1, gr) {
      fornum(j, 0, abv[mk[i]].size()) {
        ll b = abv[mk[i]][j];
        if (mk[i + 1] == b)
          continue;
        if (mk[i - 1] == b)
          continue;
        // printf("%lld,%lld,%lld:\\n", mk[0], mk[d], mk[i]);
        ll ra, rb;
        dfs2(b, mk[i], 1, i, ra, rb);
        if (gr == i * 2) {
          ret *= ra + rb * 2;
        } else {
          ret *= ra * 2 + rb;
        }
        ret %= MOD;
      }
    }
    return ret;
  }
  fornum(i, 0, abv[a].size()) {
    ll b = abv[a][i];
    if (b == p)
      continue;
    ret += dfs1(b, a, d + 1);
    ret %= MOD;
  }
  return ret;
}
int main() {
  scanf("%lld", &N);
  fornum(i, 1, N) {
    ll a, b;
    scanf("%lld%lld", &a, &b);
    abv[a].push_back(b);
    abv[b].push_back(a);
  }
  gr = grcalc(1, 0).first;
  ans = 0;
  p3[0] = 1;
  fornum(i, 0, N) { p3[i + 1] = p3[i] * 3 % MOD; }
  // printf("%lld ", gr);
  fornum(i, 0, N) {
    if (abv[i + 1].size() == 1) {
      // printf("%lld %lld\\n", i + 1, ans);
      ans += dfs1(i + 1, 0, 0);
      ans %= MOD;
    }
  }
  printf("%lld", ans);

  return 0;
}