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

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

typedef struct {
  int to, did;
} Edge;

int N;
vector<vector<Edge>> G;
vector<int> P, D;
vector<bool> done;
vector<pair<int, int>> par;
vector<bool> is_leaf;

// weight, idx
pair<int, int> dfs(int u, int p) {
  pair<int, int> r(0, u);
  for (int i = 0; i < G[u].size(); i++) {
    auto e = G[u][i];
    if (e.to == p || is_leaf[e.to])
      continue;

    par[e.to] = make_pair(u, i);
    auto t = dfs(e.to, u);
    t.first += D[e.did];
    if (r.first < t.first)
      r = t;
  }
  return r;
}
pair<int, int> diameter() {
  par.clear();
  par.resize(N, {-1, -1});
  for (int i = 0; i < N; i++)
    if (!is_leaf[i]) {
      auto r = dfs(i, -1);
      auto t = dfs(r.second, -1);
      return make_pair(r.second, t.second);
    }
  // Not Implemented
  return pair<int, int>();
}

int main() {
  while (cin >> N, N) {
    P.clear();
    P.resize(N - 1);
    for (auto &v : P)
      cin >> v;
    D.clear();
    D.resize(N - 1);
    for (auto &v : D)
      cin >> v;
    G.clear();
    G.resize(N);
    for (int i = 0; i < N - 1; i++) {
      G[i + 1].push_back({P[i] - 1, i});
      G[P[i] - 1].push_back({i + 1, i});
    }

    int ans = 0;
    is_leaf.clear();
    is_leaf.resize(N, false);
    done.clear();
    done.resize(N - 1, false);
    for (int i = 0; i < N; i++) {
      if (G[i].size() == 1) {
        ans += D[G[i][0].did];
        done[G[i][0].did] = true;
        is_leaf[i] = true;
      }
    }
    auto path = diameter();
    int v = path.second;
    while (true) {
      auto p = par[v];
      if (p.first < 0)
        break;

      int did = G[p.first][p.second].did;
      done[did] = true;
      ans += D[did] * 2;
      v = p.first;
      if (v == path.first)
        break;
    }
    for (int i = 0; i < N - 1; i++)
      if (!done[i])
        ans += D[i] * 3;
    cout << ans << endl;
  }

  return 0;
}

Error: TLE

Fixed submission:

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

typedef struct {
  int to, did;
} Edge;

int N;
vector<vector<Edge>> G;
vector<int> P, D;
vector<bool> done;
vector<pair<int, int>> par;
vector<bool> is_leaf;

// weight, idx
pair<int, int> dfs(int u, int p) {
  pair<int, int> r(0, u);
  for (int i = 0; i < G[u].size(); i++) {
    auto e = G[u][i];
    if (e.to == p || is_leaf[e.to])
      continue;

    par[e.to] = make_pair(u, i);
    auto t = dfs(e.to, u);
    t.first += D[e.did];
    if (r.first < t.first)
      r = t;
  }
  return r;
}
pair<int, int> diameter() {
  par.resize(N, {-1, -1});
  for (int i = 0; i < N; i++)
    if (!is_leaf[i]) {
      auto r = dfs(i, -1);
      auto t = dfs(r.second, -1);
      return make_pair(r.second, t.second);
    }
  // Not Implemented
  return pair<int, int>();
}

int main() {
  while (cin >> N) {
    P.clear();
    P.resize(N - 1);
    for (auto &v : P)
      cin >> v;
    D.clear();
    D.resize(N - 1);
    for (auto &v : D)
      cin >> v;
    G.clear();
    G.resize(N);
    for (int i = 0; i < N - 1; i++) {
      G[i + 1].push_back({P[i] - 1, i});
      G[P[i] - 1].push_back({i + 1, i});
    }

    int ans = 0;
    is_leaf.clear();
    is_leaf.resize(N, false);
    done.clear();
    done.resize(N - 1, false);
    for (int i = 0; i < N; i++) {
      if (G[i].size() == 1) {
        ans += D[G[i][0].did];
        done[G[i][0].did] = true;
        is_leaf[i] = true;
      }
    }
    auto path = diameter();
    int v = path.second;
    while (true) {
      auto p = par[v];
      if (p.first < 0)
        break;

      int did = G[p.first][p.second].did;
      done[did] = true;
      ans += D[did] * 2;
      v = p.first;
      if (v == path.first)
        break;
    }
    for (int i = 0; i < N - 1; i++)
      if (!done[i])
        ans += D[i] * 3;
    cout << ans << endl;
  }

  return 0;
}