Problem p03794 (C++)

Buggy submission (status: Runtime Error):
#ifdef DEBUG
#define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;

mt19937 mrand(random_device{}());

int rnd(int x) { return mrand() % x; }

typedef long double ld;
typedef long long ll;

#ifdef DEBUG
#define eprintf(...) fprintf(stderr, __VA_ARGS__), fflush(stderr)
#else
#define eprintf(...) ;
#endif

#define pb push_back
#define mp make_pair
#define sz(x) ((int)(x).size())
#define TASK "text"

const int inf = (int)1.01e9;
const ld eps = 1e-9;
const ld pi = acos((ld)-1.0);

const int mod = (int)1e9 + 7;

void add(int &x, int y) {
  if ((x += y) >= mod) {
    x -= mod;
  }
}

int mult(int x, int y) { return (long long)x * y % mod; }

int myPower(int x, int pw) {
  int res = 1;
  for (; pw; pw >>= 1) {
    if (pw & 1) {
      res = mult(res, x);
    }
    x = mult(x, x);
  }
  return res;
}

void precalc() {}

const int maxn = (int)1e3 + 10;
int n;

vector<vector<int>> es;

int read() {
  if (scanf("%d", &n) < 1) {
    return 0;
  }
  es = vector<vector<int>>(n);
  for (int i = 0; i < n - 1; ++i) {
    int s, t;
    scanf("%d%d", &s, &t);
    --s, --t;
    es[s].pb(t), es[t].pb(s);
  }
  return 1;
}

int dist[maxn];
int dist1[maxn];

void dfs1(int v, int p) {
  for (int u : es[v]) {
    if (u == p) {
      continue;
    }
    dist[u] = dist[v] + 1;
    dfs1(u, v);
  }
}

int used[maxn];

typedef vector<vector<int>> vvi;

int cnt[maxn];

void dfs2(int v, int p) {
  cnt[v] = 1;
  for (int u : es[v]) {
    if (u == p || used[u]) {
      continue;
    }
    dfs2(u, v);
    cnt[v] += cnt[u];
  }
}

int diam, D;

void adde(vvi &a) {
  int n = sz(a);
  for (int i = 0; i < n; ++i) {
    a[i].pb(0);
  }
  a.pb(vector<int>(n + 1, 0));

  for (int i = n - 1; i >= 0; --i) {
    for (int j = n - 1; j >= 0; --j) {
      int x = a[i][j];
      if (!x) {
        continue;
      }
      a[i][j] = 0;
      if (i + 1 <= D) {
        add(a[i + 1][j], x);
      }
      if (i + 1 <= D) {
        add(a[i][j + 1], x);
      }
    }
  }
}

void makesum(vvi &a) {
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      int &res = a[i][j];
      if (i) {
        add(res, a[i - 1][j]);
        if (j) {
          add(res, mod - a[i - 1][j - 1]);
        }
      }
      if (j) {
        add(res, a[i][j - 1]);
      }
    }
  }
}

void makeone(vvi &a) {
  for (int i = sz(a) - 1; i >= 0; --i) {
    for (int j = sz(a) - 1; j >= 0; --j) {
      int &res = a[i][j];
      if (i) {
        add(res, mod - a[i - 1][j]);
        if (j) {
          add(res, a[i - 1][j - 1]);
        }
      }
      if (j) {
        add(res, mod - a[i][j - 1]);
      }
    }
  }
}

void update(vvi &a, vvi &b) {
  /*eprintf("update:\n");
  eprintf("%d\n", sz(a));
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      eprintf("%d%c", a[i][j], " \n"[j == sz(a) - 1]);
    }
  }
  eprintf("%d\n", sz(b));
  for (int i = 0; i < sz(b); ++i) {
    for (int j = 0; j < sz(b); ++j) {
      eprintf("%d%c", b[i][j], " \n"[j == sz(b) - 1]);
    }
  }*/
  assert(sz(a) >= sz(b));

  makesum(a);
  makesum(b);
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      a[i][j] = mult(a[i][j], b[min(sz(b) - 1, i)][min(sz(b) - 1, j)]);
    }
  }
  makeone(a);
}

vvi *dfs3(int v, int p) {
  int u0 = -1;
  for (int u : es[v]) {
    if (u == p || used[u]) {
      continue;
    }
    if (u0 == -1 || cnt[u] > cnt[u0]) {
      u0 = u;
    }
  }
  if (u0 == -1) {
    return new vvi({{1}});
  }
  auto ans = dfs3(u0, v);
  adde(*ans);

  for (int u : es[v]) {
    if (u == p || u == u0 || used[u]) {
      continue;
    }
    auto got = dfs3(u, v);
    adde(*got);
    update(*ans, *got);
  }
  return ans;
}

int ans[maxn];
int adds[maxn];

void solve() {
  dist[0] = 0;
  dfs1(0, -1);
  int v = max_element(dist, dist + n) - dist;

  dist[v] = 0;
  dfs1(v, -1);
  int u = max_element(dist, dist + n) - dist;
  diam = dist[u];

  memcpy(dist1, dist, sizeof(dist[0]) * n);

  dist[u] = 0;
  dfs1(u, -1);

  D = (diam + 1) / 2;
  // eprintf("diam = %d\n", diam);

  for (int i = 0; i < n; ++i) {
    used[i] = 0;
  }
  int vs[2];
  int k = 0;
  for (int i = 0; i < n; ++i) {
    if (dist[i] + dist1[i] == diam && abs(dist[i] - dist1[i]) <= 1) {
      vs[k++] = i;
      used[i] = 1;
    }
  }
  assert(1 <= k && k <= 2);
  for (int i = 0; i < k; ++i) {
    dfs2(i, -1);
  }

  // eprintf("middle %d\n", vs[0]);

  int res = 0;
  if (k == 1) {
    dfs2(vs[0], -1);
    for (int i = 0; i <= D; ++i) {
      ans[i] = 1;
    }
    for (int u : es[vs[0]]) {
      auto &got = *dfs3(u, vs[0]);
      adde(got);
      for (int i = 0; i <= D; ++i) {
        adds[i] = 0;
      }
      for (int i = 0; i < sz(got); ++i) {
        for (int j = 0; j < sz(got); ++j) {
          int cur = got[i][j];
          if (!cur) {
            continue;
          }
          // eprintf("u = %d, (%d,%d) x%d\n", u, i, j, cur);
          int l = i, r = D - j;
          if (l <= r) {
            // eprintf("[%d..%d]\n", l, r);
            add(adds[l], cur);
            add(adds[r + 1], mod - cur);
          }
        }
      }
      int bal = 0;
      for (int i = 0; i <= D; ++i) {
        add(bal, adds[i]);
        ans[i] = mult(ans[i], bal);
      }
    }
    for (int i = 0; i <= D; ++i) {
      add(res, ans[i]);
    }
  } else {
    auto &ans1 = *(dfs3(vs[0], -1));
    auto &ans2 = *(dfs3(vs[1], -1));
    assert(sz(ans1) == D && sz(ans2) == sz(ans1));
    // eprintf("D = %d\n", D);
    for (int i = 0; i < sz(ans1); ++i) {
      add(res, mult(2, mult(ans1[i][D - 1 - i], ans2[i][D - 1 - i])));
      if (i) {
        add(res, mult(ans1[i][D - 1 - i], ans2[i][D - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i][D - i]));

        add(res, mult(ans1[i][D - 1 - i], ans2[i - 1][D - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i - 1][D - i]));
      }
      if (i < sz(ans1) - 1) {
        add(res, mult(ans1[i][D - 1 - i], ans2[i + 1][D - 1 - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i + 1][D - 1 - i]));
      }
    }
  }

  printf("%d\n", res);
  // exit(0);
}

int main() {
  precalc();
#ifdef LOCAL
  freopen(TASK ".out", "w", stdout);
  assert(freopen(TASK ".in", "r", stdin));
#endif

  while (1) {
    if (!read()) {
      break;
    }
    solve();
#ifdef DEBUG
    eprintf("Time %.2f\n", (double)clock() / CLOCKS_PER_SEC);
#endif
  }
  return 0;
}


Fixed submission:
#ifdef DEBUG
#define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>

using namespace std;

mt19937 mrand(random_device{}());

int rnd(int x) { return mrand() % x; }

typedef long double ld;
typedef long long ll;

#ifdef DEBUG
#define eprintf(...) fprintf(stderr, __VA_ARGS__), fflush(stderr)
#else
#define eprintf(...) ;
#endif

#define pb push_back
#define mp make_pair
#define sz(x) ((int)(x).size())
#define TASK "text"

const int inf = (int)1.01e9;
const ld eps = 1e-9;
const ld pi = acos((ld)-1.0);

const int mod = (int)1e9 + 7;

void add(int &x, int y) {
  if ((x += y) >= mod) {
    x -= mod;
  }
}

int mult(int x, int y) { return (long long)x * y % mod; }

int myPower(int x, int pw) {
  int res = 1;
  for (; pw; pw >>= 1) {
    if (pw & 1) {
      res = mult(res, x);
    }
    x = mult(x, x);
  }
  return res;
}

void precalc() {}

const int maxn = (int)1e3 + 10;
int n;

vector<vector<int>> es;

int read() {
  if (scanf("%d", &n) < 1) {
    return 0;
  }
  es = vector<vector<int>>(n);
  for (int i = 0; i < n - 1; ++i) {
    int s, t;
    scanf("%d%d", &s, &t);
    --s, --t;
    es[s].pb(t), es[t].pb(s);
  }
  return 1;
}

int dist[maxn];
int dist1[maxn];

void dfs1(int v, int p) {
  for (int u : es[v]) {
    if (u == p) {
      continue;
    }
    dist[u] = dist[v] + 1;
    dfs1(u, v);
  }
}

int used[maxn];

typedef vector<vector<int>> vvi;

int cnt[maxn];

void dfs2(int v, int p) {
  cnt[v] = 1;
  for (int u : es[v]) {
    if (u == p || used[u]) {
      continue;
    }
    dfs2(u, v);
    cnt[v] += cnt[u];
  }
}

int diam, D;

void adde(vvi &a) {
  int n = sz(a);
  for (int i = 0; i < n; ++i) {
    a[i].pb(0);
  }
  a.pb(vector<int>(n + 1, 0));

  for (int i = n - 1; i >= 0; --i) {
    for (int j = n - 1; j >= 0; --j) {
      int x = a[i][j];
      if (!x) {
        continue;
      }
      a[i][j] = 0;
      if (i + 1 <= D) {
        add(a[i + 1][j], x);
      }
      if (i + 1 <= D) {
        add(a[i][j + 1], x);
      }
    }
  }
}

void makesum(vvi &a) {
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      int &res = a[i][j];
      if (i) {
        add(res, a[i - 1][j]);
        if (j) {
          add(res, mod - a[i - 1][j - 1]);
        }
      }
      if (j) {
        add(res, a[i][j - 1]);
      }
    }
  }
}

void makeone(vvi &a) {
  for (int i = sz(a) - 1; i >= 0; --i) {
    for (int j = sz(a) - 1; j >= 0; --j) {
      int &res = a[i][j];
      if (i) {
        add(res, mod - a[i - 1][j]);
        if (j) {
          add(res, a[i - 1][j - 1]);
        }
      }
      if (j) {
        add(res, mod - a[i][j - 1]);
      }
    }
  }
}

void update(vvi &a, vvi &b) {
  /*eprintf("update:\n");
  eprintf("%d\n", sz(a));
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      eprintf("%d%c", a[i][j], " \n"[j == sz(a) - 1]);
    }
  }
  eprintf("%d\n", sz(b));
  for (int i = 0; i < sz(b); ++i) {
    for (int j = 0; j < sz(b); ++j) {
      eprintf("%d%c", b[i][j], " \n"[j == sz(b) - 1]);
    }
  }*/
  assert(sz(a) >= sz(b));

  makesum(a);
  makesum(b);
  for (int i = 0; i < sz(a); ++i) {
    for (int j = 0; j < sz(a); ++j) {
      a[i][j] = mult(a[i][j], b[min(sz(b) - 1, i)][min(sz(b) - 1, j)]);
    }
  }
  makeone(a);
}

vvi *dfs3(int v, int p) {
  int u0 = -1;
  for (int u : es[v]) {
    if (u == p || used[u]) {
      continue;
    }
    if (u0 == -1 || cnt[u] > cnt[u0]) {
      u0 = u;
    }
  }
  if (u0 == -1) {
    return new vvi({{1}});
  }
  auto ans = dfs3(u0, v);
  adde(*ans);

  for (int u : es[v]) {
    if (u == p || u == u0 || used[u]) {
      continue;
    }
    auto got = dfs3(u, v);
    adde(*got);
    if (sz(*ans) < sz(*got)) {
      swap(ans, got);
    }
    update(*ans, *got);
  }
  return ans;
}

int ans[maxn];
int adds[maxn];

void solve() {
  dist[0] = 0;
  dfs1(0, -1);
  int v = max_element(dist, dist + n) - dist;

  dist[v] = 0;
  dfs1(v, -1);
  int u = max_element(dist, dist + n) - dist;
  diam = dist[u];

  memcpy(dist1, dist, sizeof(dist[0]) * n);

  dist[u] = 0;
  dfs1(u, -1);

  D = (diam + 1) / 2;
  // eprintf("diam = %d\n", diam);

  for (int i = 0; i < n; ++i) {
    used[i] = 0;
  }
  int vs[2];
  int k = 0;
  for (int i = 0; i < n; ++i) {
    if (dist[i] + dist1[i] == diam && abs(dist[i] - dist1[i]) <= 1) {
      vs[k++] = i;
      used[i] = 1;
    }
  }
  assert(1 <= k && k <= 2);
  for (int i = 0; i < k; ++i) {
    dfs2(i, -1);
  }

  // eprintf("middle %d\n", vs[0]);

  int res = 0;
  if (k == 1) {
    dfs2(vs[0], -1);
    for (int i = 0; i <= D; ++i) {
      ans[i] = 1;
    }
    for (int u : es[vs[0]]) {
      auto &got = *dfs3(u, vs[0]);
      adde(got);
      for (int i = 0; i <= D; ++i) {
        adds[i] = 0;
      }
      for (int i = 0; i < sz(got); ++i) {
        for (int j = 0; j < sz(got); ++j) {
          int cur = got[i][j];
          if (!cur) {
            continue;
          }
          // eprintf("u = %d, (%d,%d) x%d\n", u, i, j, cur);
          int l = i, r = D - j;
          if (l <= r) {
            // eprintf("[%d..%d]\n", l, r);
            add(adds[l], cur);
            add(adds[r + 1], mod - cur);
          }
        }
      }
      int bal = 0;
      for (int i = 0; i <= D; ++i) {
        add(bal, adds[i]);
        ans[i] = mult(ans[i], bal);
      }
    }
    for (int i = 0; i <= D; ++i) {
      add(res, ans[i]);
    }
  } else {
    auto &ans1 = *(dfs3(vs[0], -1));
    auto &ans2 = *(dfs3(vs[1], -1));
    assert(sz(ans1) == D && sz(ans2) == sz(ans1));
    // eprintf("D = %d\n", D);
    for (int i = 0; i < sz(ans1); ++i) {
      add(res, mult(2, mult(ans1[i][D - 1 - i], ans2[i][D - 1 - i])));
      if (i) {
        add(res, mult(ans1[i][D - 1 - i], ans2[i][D - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i][D - i]));

        add(res, mult(ans1[i][D - 1 - i], ans2[i - 1][D - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i - 1][D - i]));
      }
      if (i < sz(ans1) - 1) {
        add(res, mult(ans1[i][D - 1 - i], ans2[i + 1][D - 1 - i]));
        add(res, mult(ans2[i][D - 1 - i], ans1[i + 1][D - 1 - i]));
      }
    }
  }

  printf("%d\n", res);
  // exit(0);
}

int main() {
  precalc();
#ifdef LOCAL
  freopen(TASK ".out", "w", stdout);
  assert(freopen(TASK ".in", "r", stdin));
#endif

  while (1) {
    if (!read()) {
      break;
    }
    solve();
#ifdef DEBUG
    eprintf("Time %.2f\n", (double)clock() / CLOCKS_PER_SEC);
#endif
  }
  return 0;
}