Problem p03384 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
template <typename A, typename B>
ostream &operator<<(ostream &cout, pair<A, B> const &p) {
  return cout << "(" << p.first << ", " << p.second << ")";
}
template <typename A> ostream &operator<<(ostream &cout, vector<A> const &v) {
  cout << "{";
  for (int i = 0; i < v.size(); i++) {
    if (i)
      cout << ", ";
    cout << v[i];
  }
  return cout << "}";
}

int main() {
  //    ios_base::sync_with_stdio(0); cin.tie(0);
  freopen("file.in", "r", stdin);
  int n;
  cin >> n;
  vector<vector<int>> adj(n);
  vector<int> depth(n);
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--;
    v--;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  pair<int, pair<int, int>> diam = {0, {0, 0}};
  function<pair<int, int>(int, int)> dfs_diam = [&](int src,
                                                    int par) -> pair<int, int> {
    pair<int, int> mx = {depth[src], src};
    set<pair<int, int>, greater<pair<int, int>>> store;
    for (int nxt : adj[src]) {
      if (nxt == par)
        continue;
      depth[nxt] = depth[src] + 1;
      auto val = dfs_diam(nxt, src);
      mx = max(mx, val);
      store.insert(val);
    }
    if (mx.first > diam.first) {
      diam = {mx.first, {0, mx.second}};
    }
    if (store.size() >= 2) {
      int cnt = 0;
      pair<int, int> one = *store.begin();
      pair<int, int> two = *next(store.begin());
      if (one.first + two.first - 2 * depth[src] > diam.first) {
        diam.first = one.first + two.first;
        diam.second.first = one.second;
        diam.second.second = two.second;
      }
    }
    return mx;
  };
  dfs_diam(0, -1);
  vector<int> path;
  vector<bool> on_path(n);
  int st = diam.second.first;
  int ed = diam.second.second;
  function<bool(int, int)> dfs_path = [&](int src, int par) -> bool {
    path.push_back(src);
    if (src == ed) {
      for (int x : path)
        on_path[x] = true;
      return true;
    }
    for (int nxt : adj[src]) {
      if (nxt == par)
        continue;
      if (dfs_path(nxt, src))
        return true;
    }
    path.pop_back();
    return false;
  };
  dfs_path(st, -1);
  for (int i = 0; i < n; i++) {
    if (!on_path[i] && adj[i].size() > 1) {
      cout << -1 << '\n';
      return 0;
    }
  }
  vector<int> sizes;
  for (int x : path) {
    if (adj[x].size() == 1) {
      sizes.push_back(0);
    } else {
      sizes.push_back(adj[x].size() - 2);
    }
  }
  auto solve = [&](vector<int> use) -> vector<int> {
    int sz = use.size();
    vector<int> res;
    res.push_back(1);
    int cur = 1;
    for (int i = 1; i < sz; i++) {
      cur += (use[i - 1] + 1);
      res.push_back(cur);
    }
    vector<int> fin;
    fin.push_back(1);
    for (int i = 1; i < res.size(); i++) {
      if (i != res.size() - 1) {
        for (int j = res[i] + 1; j < res[i + 1]; j++) {
          fin.push_back(j);
        }
      }
      fin.push_back(res[i]);
    }
    return fin;
  };
  vector<int> c1 = solve(sizes);
  reverse(sizes.begin(), sizes.end());
  vector<int> c2 = solve(sizes);
  vector<int> ans = c1;
  for (int i = 0; i < n; i++) {
    if (c1[i] < c2[i]) {
      ans = c1;
      break;
    } else if (c1[i] > c2[i]) {
      ans = c2;
      break;
    }
  }
  for (int i = 0; i < n; i++) {
    cout << ans[i] << " ";
  }
  cout << '\n';
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
template <typename A, typename B>
ostream &operator<<(ostream &cout, pair<A, B> const &p) {
  return cout << "(" << p.first << ", " << p.second << ")";
}
template <typename A> ostream &operator<<(ostream &cout, vector<A> const &v) {
  cout << "{";
  for (int i = 0; i < v.size(); i++) {
    if (i)
      cout << ", ";
    cout << v[i];
  }
  return cout << "}";
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  int n;
  cin >> n;
  vector<vector<int>> adj(n);
  vector<int> depth(n);
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--;
    v--;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  pair<int, pair<int, int>> diam = {0, {0, 0}};
  function<pair<int, int>(int, int)> dfs_diam = [&](int src,
                                                    int par) -> pair<int, int> {
    pair<int, int> mx = {depth[src], src};
    set<pair<int, int>, greater<pair<int, int>>> store;
    for (int nxt : adj[src]) {
      if (nxt == par)
        continue;
      depth[nxt] = depth[src] + 1;
      auto val = dfs_diam(nxt, src);
      mx = max(mx, val);
      store.insert(val);
    }
    if (mx.first > diam.first) {
      diam = {mx.first, {0, mx.second}};
    }
    if (store.size() >= 2) {
      int cnt = 0;
      pair<int, int> one = *store.begin();
      pair<int, int> two = *next(store.begin());
      if (one.first + two.first - 2 * depth[src] > diam.first) {
        diam.first = one.first + two.first;
        diam.second.first = one.second;
        diam.second.second = two.second;
      }
    }
    return mx;
  };
  dfs_diam(0, -1);
  vector<int> path;
  vector<bool> on_path(n);
  int st = diam.second.first;
  int ed = diam.second.second;
  function<bool(int, int)> dfs_path = [&](int src, int par) -> bool {
    path.push_back(src);
    if (src == ed) {
      for (int x : path)
        on_path[x] = true;
      return true;
    }
    for (int nxt : adj[src]) {
      if (nxt == par)
        continue;
      if (dfs_path(nxt, src))
        return true;
    }
    path.pop_back();
    return false;
  };
  dfs_path(st, -1);
  for (int i = 0; i < n; i++) {
    if (!on_path[i] && adj[i].size() > 1) {
      cout << -1 << '\n';
      return 0;
    }
  }
  vector<int> sizes;
  for (int x : path) {
    if (adj[x].size() == 1) {
      sizes.push_back(0);
    } else {
      sizes.push_back(adj[x].size() - 2);
    }
  }
  auto solve = [&](vector<int> use) -> vector<int> {
    int sz = use.size();
    vector<int> res;
    res.push_back(1);
    int cur = 1;
    for (int i = 1; i < sz; i++) {
      cur += (use[i - 1] + 1);
      res.push_back(cur);
    }
    vector<int> fin;
    fin.push_back(1);
    for (int i = 1; i < res.size(); i++) {
      if (i != res.size() - 1) {
        for (int j = res[i] + 1; j < res[i + 1]; j++) {
          fin.push_back(j);
        }
      }
      fin.push_back(res[i]);
    }
    return fin;
  };
  vector<int> c1 = solve(sizes);
  reverse(sizes.begin(), sizes.end());
  vector<int> c2 = solve(sizes);
  vector<int> ans = c1;
  for (int i = 0; i < n; i++) {
    if (c1[i] < c2[i]) {
      ans = c1;
      break;
    } else if (c1[i] > c2[i]) {
      ans = c2;
      break;
    }
  }
  for (int i = 0; i < n; i++) {
    cout << ans[i] << " ";
  }
  cout << '\n';
  return 0;
}