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

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

#define int long long

#define rep(i, n) for (int i = 0; i < (n); i++)
#define pb push_back
#define all(v) (v).begin(), (v).end()
#define fi first
#define se second
typedef vector<int> vint;
typedef pair<int, int> pint;
typedef vector<pint> vpint;

template <typename A, typename B> inline void chmin(A &a, B b) {
  if (a > b)
    a = b;
}
template <typename A, typename B> inline void chmax(A &a, B b) {
  if (a < b)
    a = b;
}

const int mod = 1000000007;
inline void add(int &a, int b) {
  a += b;
  if (a >= mod)
    a -= mod;
}
int mpow(int n, int m) {
  int ret = 1;
  while (m) {
    if (m & 1)
      ret = ret * n % mod;
    n = n * n % mod;
    m >>= 1;
  }
  return ret;
}
const int FACT_SIZE = 1111111;
int fact[FACT_SIZE];
int inv[FACT_SIZE];
struct fact_exec {
  fact_exec() {
    fact[0] = 1;
    for (int i = 1; i < FACT_SIZE; i++)
      fact[i] = fact[i - 1] * i % mod;
    inv[FACT_SIZE - 1] = mpow(fact[FACT_SIZE - 1], mod - 2);
    for (int i = FACT_SIZE - 2; i >= 0; i--)
      inv[i] = inv[i + 1] * (i + 1) % mod;
  }
} factexec;
int nCk(int n, int k) {
  if (n < 0 | k < 0 || k > n)
    return 0;
  return fact[n] * inv[k] % mod * inv[n - k] % mod;
}
int nPk(int n, int k) {
  if (n < 0 || k < 0 || k > n)
    return 0;
  return fact[n] * inv[n - k] % mod;
}

int N;
vint G[1111];

int dist[1111];

void dfs(int v, int p, int d) {
  dist[v] = d;
  for (auto u : G[v]) {
    if (u == p)
      continue;
    dfs(u, v, d + 1);
  }
}

int D;
int X;
int dp[1111][1111];
void dfs2(int v, int p, int d) {

  int l = max(0ll, X + d - D / 2);

  for (int i = l; i <= X; i++)
    dp[v][i] = 1;
  for (auto u : G[v]) {
    if (u == p)
      continue;
    dfs2(u, v, d + 1);
    for (int i = l; i <= X; i++)
      dp[v][i] = dp[v][i] * (dp[u][i] + dp[u][i + 1]) % mod;
  }
}
void solve(int c) {
  int ans = 0;
  for (X = 0; X <= D / 2; X++) {
    memset(dp, 0, sizeof(dp));
    dfs2(c, -1, 0);
    add(ans, dp[c][0]);
  }
  cout << ans << endl;
}

signed main() {

  cin >> N;
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    G[a].pb(b);
    G[b].pb(a);
  }

  dfs(0, -1, 0);
  int v = 0;
  rep(i, N) if (dist[i] > dist[v]) v = i;
  dfs(v, -1, 0);

  int u = 0;
  rep(i, N) if (dist[i] > dist[u]) u = i;

  vint latte(N);
  rep(i, N) latte[i] = dist[i];
  dfs(u, -1, 0);
  rep(i, N) chmax(latte[i], dist[i]);

  vint lis;
  rep(i, N) if (latte[i] <= (dist[v] + 1) / 2) lis.pb(i);

  D = dist[v];

  if (lis.size() == 1) {
    solve(lis[0]);
  } else {
    int u = lis[0], v = lis[1];

    vint x(N), y(N), z(N), w(N);

    for (X = 0; X <= D / 2; X++) {
      memset(dp, 0, sizeof(dp));
      dfs2(u, v, 0);
      x[X] = dp[u][0];

      memset(dp, 0, sizeof(dp));
      dfs2(v, u, 0);
      y[X] = dp[v][0];
    }

    D++;
    for (X = 0; X <= D / 2; X++) {
      memset(dp, 0, sizeof(dp));
      dfs2(u, v, 0);
      z[X] = dp[u][0];

      memset(dp, 0, sizeof(dp));
      dfs2(v, u, 0);
      w[X] = dp[v][0];
    }

    D--;
    z[0] = w[0] = 0;
    for (int i = 1; i < N; i++)
      z[i] = (z[i] - x[i] - x[i - 1] + mod * 2) % mod;
    for (int i = 1; i < N; i++)
      w[i] = (w[i] - y[i] - y[i - 1] + mod * 2) % mod;

    int ans = 0;
    for (int i = 0; i <= D / 2; i++) {
      for (int j = 0; j <= D / 2; j++) {
        int a = i;
        int b = D / 2 - i;
        int c = j;
        int d = D / 2 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, x[i] * y[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, x[i] * y[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2; i++) {
      for (int j = 0; j <= D / 2 + 1; j++) {
        int a = i;
        int b = D / 2 - i;
        int c = j;
        int d = D / 2 + 1 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, x[i] * w[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, x[i] * w[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2 + 1; i++) {
      for (int j = 0; j <= D / 2 + 1; j++) {
        int a = i;
        int b = D / 2 + 1 - i;
        int c = j;
        int d = D / 2 + 1 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, z[i] * w[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, z[i] * w[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2 + 1; i++) {
      for (int j = 0; j <= D / 2; j++) {
        int a = i;
        int b = D / 2 + 1 - i;
        int c = j;
        int d = D / 2 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, z[i] * y[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, z[i] * y[j] % mod);
      }
    }

    cout << ans << endl;
  }
  return 0;
}


Error: TLE

Fixed submission:

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

#define int long long

#define rep(i, n) for (int i = 0; i < (n); i++)
#define pb push_back
#define all(v) (v).begin(), (v).end()
#define fi first
#define se second
typedef vector<int> vint;
typedef pair<int, int> pint;
typedef vector<pint> vpint;

template <typename A, typename B> inline void chmin(A &a, B b) {
  if (a > b)
    a = b;
}
template <typename A, typename B> inline void chmax(A &a, B b) {
  if (a < b)
    a = b;
}

const int mod = 1000000007;
inline void add(int &a, int b) {
  a += b;
  if (a >= mod)
    a -= mod;
}
int mpow(int n, int m) {
  int ret = 1;
  while (m) {
    if (m & 1)
      ret = ret * n % mod;
    n = n * n % mod;
    m >>= 1;
  }
  return ret;
}
const int FACT_SIZE = 1111111;
int fact[FACT_SIZE];
int inv[FACT_SIZE];
struct fact_exec {
  fact_exec() {
    fact[0] = 1;
    for (int i = 1; i < FACT_SIZE; i++)
      fact[i] = fact[i - 1] * i % mod;
    inv[FACT_SIZE - 1] = mpow(fact[FACT_SIZE - 1], mod - 2);
    for (int i = FACT_SIZE - 2; i >= 0; i--)
      inv[i] = inv[i + 1] * (i + 1) % mod;
  }
} factexec;
int nCk(int n, int k) {
  if (n < 0 | k < 0 || k > n)
    return 0;
  return fact[n] * inv[k] % mod * inv[n - k] % mod;
}
int nPk(int n, int k) {
  if (n < 0 || k < 0 || k > n)
    return 0;
  return fact[n] * inv[n - k] % mod;
}

int N;
vint G[1111];

int dist[1111];

void dfs(int v, int p, int d) {
  dist[v] = d;
  for (auto u : G[v]) {
    if (u == p)
      continue;
    dfs(u, v, d + 1);
  }
}

int D;
int X;
int dp[1111][1111];
void dfs2(int v, int p, int d) {
  rep(i, D / 2 + 1) dp[v][i] = 1;
  for (auto u : G[v]) {
    if (u == p)
      continue;
    dfs2(u, v, d + 1);
    rep(i, D / 2 + 1) dp[v][i] = dp[v][i] * (dp[u][i] + dp[u][i + 1]) % mod;
  }
  rep(i, D / 2 + 1) if (i > X || d - i > D / 2 - X) dp[v][i] = 0;
}
void solve(int c) {
  int ans = 0;
  for (X = 0; X <= D / 2; X++) {
    memset(dp, 0, sizeof(dp));
    dfs2(c, -1, 0);
    add(ans, dp[c][0]);
  }
  cout << ans << endl;
}

signed main() {

  cin >> N;
  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    G[a].pb(b);
    G[b].pb(a);
  }

  dfs(0, -1, 0);
  int v = 0;
  rep(i, N) if (dist[i] > dist[v]) v = i;
  dfs(v, -1, 0);

  int u = 0;
  rep(i, N) if (dist[i] > dist[u]) u = i;

  vint latte(N);
  rep(i, N) latte[i] = dist[i];
  dfs(u, -1, 0);
  rep(i, N) chmax(latte[i], dist[i]);

  vint lis;
  rep(i, N) if (latte[i] <= (dist[v] + 1) / 2) lis.pb(i);

  D = dist[v];

  if (lis.size() == 1) {
    solve(lis[0]);
  } else {
    int u = lis[0], v = lis[1];

    vint x(N), y(N), z(N), w(N);

    for (X = 0; X <= D / 2; X++) {
      memset(dp, 0, sizeof(dp));
      dfs2(u, v, 0);
      x[X] = dp[u][0];

      memset(dp, 0, sizeof(dp));
      dfs2(v, u, 0);
      y[X] = dp[v][0];
    }

    D++;
    for (X = 0; X <= D / 2; X++) {
      memset(dp, 0, sizeof(dp));
      dfs2(u, v, 0);
      z[X] = dp[u][0];

      memset(dp, 0, sizeof(dp));
      dfs2(v, u, 0);
      w[X] = dp[v][0];
    }

    D--;
    z[0] = w[0] = 0;
    for (int i = 1; i < N; i++)
      z[i] = (z[i] - x[i] - x[i - 1] + mod * 2) % mod;
    for (int i = 1; i < N; i++)
      w[i] = (w[i] - y[i] - y[i - 1] + mod * 2) % mod;

    int ans = 0;
    for (int i = 0; i <= D / 2; i++) {
      for (int j = 0; j <= D / 2; j++) {
        int a = i;
        int b = D / 2 - i;
        int c = j;
        int d = D / 2 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, x[i] * y[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, x[i] * y[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2; i++) {
      for (int j = 0; j <= D / 2 + 1; j++) {
        int a = i;
        int b = D / 2 - i;
        int c = j;
        int d = D / 2 + 1 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, x[i] * w[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, x[i] * w[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2 + 1; i++) {
      for (int j = 0; j <= D / 2 + 1; j++) {
        int a = i;
        int b = D / 2 + 1 - i;
        int c = j;
        int d = D / 2 + 1 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, z[i] * w[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, z[i] * w[j] % mod);
      }
    }

    for (int i = 0; i <= D / 2 + 1; i++) {
      for (int j = 0; j <= D / 2; j++) {
        int a = i;
        int b = D / 2 + 1 - i;
        int c = j;
        int d = D / 2 - j;
        if (a + d <= D / 2 + 1 && b + c + 1 <= D / 2 + 1)
          add(ans, z[i] * y[j] % mod);
        if (a + d + 1 <= D / 2 + 1 && b + c <= D / 2 + 1)
          add(ans, z[i] * y[j] % mod);
      }
    }

    cout << ans << endl;
  }
  return 0;
}