Problem p03654 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int INF = (int)1.01e9;
const int K = 17;
const int N = 1 << K;

int n;
vector<pair<int, int>> e[N];
int sz[N];
ll sumlen[N];

void dfs1(int v, int pr) {
  sz[v] = 1;
  sumlen[v] = 0;
  for (auto o : e[v]) {
    int to = o.first;
    if (to == pr)
      continue;
    dfs1(to, v);
    sz[v] += sz[to];
    sumlen[v] += sumlen[to] + 1LL * o.second * sz[to];
  }
}

ll ans;

void dfs2(int v, int pr) {
  bool ok = 1;
  for (auto o : e[v]) {
    int to = o.first;
    int s1 = sz[to];
    int s2 = n - sz[to];
    if (s1 > s2 + 1)
      ok = 0;
  }

  if (ok) {
    int mn = INF;
    ll csum = 0;
    for (auto o : e[v]) {
      mn = min(mn, o.second);
      int to = o.first;
      csum += sumlen[to] + 1LL * sz[to] * o.second;
    }
    int vert = -1;
    for (auto o : e[v]) {
      int to = o.first;
      if (sz[to] >= n - sz[to]) {
        // vert = to;
        mn = o.second;
        if (sz[to] > n - sz[to])
          vert = to;
      }
    }
    if (vert != -1) {
      int cmn = INF;
      for (auto o : e[vert]) {
        if (o.first != v) {
          cmn = min(cmn, o.second);
        }
      }
      mn = 2 * mn + cmn;
    }
    ans = max(ans, 2 * csum - mn);
  }

  int m = e[v].size();
  vector<ll> pref(m + 2), suf(m + 2);
  for (int i = 1; i <= m; i++) {
    int to = e[v][i - 1].first;
    pref[i] = pref[i - 1] + sumlen[to] + 1LL * sz[to] * e[v][i - 1].second;
  }
  for (int i = m; i >= 0; i--) {
    int to = e[v][i - 1].first;
    suf[i] = suf[i + 1] + sumlen[to] + 1LL * sz[to] * e[v][i - 1].second;
  }

  for (int i = 1; i <= m; i++) {
    int to = e[v][i - 1].first;
    if (to == pr)
      continue;

    int oldSz = sz[v];
    ll oldSumlen = sumlen[v];

    sz[v] = n - sz[to];
    sumlen[v] = pref[i - 1] + suf[i + 1];
    dfs2(to, v);

    sz[v] = oldSz;
    sumlen[v] = oldSumlen;
  }
}

int main() {
#ifdef HOME
  assert(freopen("in", "r", stdin));
#endif

  while (scanf("%d", &n) == 1) {
    for (int i = 0; i < n; i++) {
      e[i].clear();
    }

    for (int i = 0; i < n - 1; i++) {
      int u, v, w;
      scanf("%d%d%d", &u, &v, &w);
      --u;
      --v;
      e[u].push_back({v, w});
      e[v].push_back({u, w});
    }

    dfs1(0, -1);
    ans = -1;
    dfs2(0, -1);
    printf("%lld\n", ans);
  }

#ifdef HOME
  cerr << "time: " << clock() * 1.0 / CLOCKS_PER_SEC << endl;
#endif
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int INF = (int)1.01e9;
const int K = 17;
const int N = 1 << K;

int n;
vector<pair<int, int>> e[N];
int sz[N];
ll sumlen[N];

void dfs1(int v, int pr) {
  sz[v] = 1;
  sumlen[v] = 0;
  for (auto o : e[v]) {
    int to = o.first;
    if (to == pr)
      continue;
    dfs1(to, v);
    sz[v] += sz[to];
    sumlen[v] += sumlen[to] + 1LL * o.second * sz[to];
  }
}

ll ans;

void dfs2(int v, int pr) {
  bool ok = 1;
  for (auto o : e[v]) {
    int to = o.first;
    int s1 = sz[to];
    int s2 = n - sz[to];
    if (s1 > s2 + 1)
      ok = 0;
  }

  if (ok) {
    int mn = INF;
    ll csum = 0;
    for (auto o : e[v]) {
      mn = min(mn, o.second);
      int to = o.first;
      csum += sumlen[to] + 1LL * sz[to] * o.second;
    }
    int vert = -1;
    for (auto o : e[v]) {
      int to = o.first;
      if (sz[to] >= n - sz[to]) {
        // vert = to;
        mn = o.second;
        if (sz[to] > n - sz[to])
          vert = to;
      }
    }
    if (vert != -1) {
      int cmn = INF;
      for (auto o : e[vert]) {
        if (o.first != v) {
          cmn = min(cmn, o.second);
        }
      }
      mn = 2 * mn + cmn;
    }
    ans = max(ans, 2 * csum - mn);
  }

  int m = e[v].size();
  vector<ll> pref(m + 2), suf(m + 2);
  for (int i = 1; i <= m; i++) {
    int to = e[v][i - 1].first;
    pref[i] = pref[i - 1] + sumlen[to] + 1LL * sz[to] * e[v][i - 1].second;
  }
  for (int i = m; i >= 1; i--) {
    int to = e[v][i - 1].first;
    suf[i] = suf[i + 1] + sumlen[to] + 1LL * sz[to] * e[v][i - 1].second;
  }

  for (int i = 1; i <= m; i++) {
    int to = e[v][i - 1].first;
    if (to == pr)
      continue;

    int oldSz = sz[v];
    ll oldSumlen = sumlen[v];

    sz[v] = n - sz[to];
    sumlen[v] = pref[i - 1] + suf[i + 1];
    dfs2(to, v);

    sz[v] = oldSz;
    sumlen[v] = oldSumlen;
  }
}

int main() {
#ifdef HOME
  assert(freopen("in", "r", stdin));
#endif

  while (scanf("%d", &n) == 1) {
    for (int i = 0; i < n; i++) {
      e[i].clear();
    }

    for (int i = 0; i < n - 1; i++) {
      int u, v, w;
      scanf("%d%d%d", &u, &v, &w);
      --u;
      --v;
      e[u].push_back({v, w});
      e[v].push_back({u, w});
    }

    dfs1(0, -1);
    ans = -1;
    dfs2(0, -1);
    printf("%lld\n", ans);
  }

#ifdef HOME
  cerr << "time: " << clock() * 1.0 / CLOCKS_PER_SEC << endl;
#endif
  return 0;
}