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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

typedef pair<int, int> pi;
typedef pair<int, pi> ppi;
#define MK make_pair
#define F first
#define S second

struct SP {
  stack<pi> s;
};

SP o[1000];

ppi search(int s, int n) {
  ppi p;
  int sum = 0;
  int mx = 0;
  int sub = 0;
  int dis = 0;
  while (!o[n].s.empty()) {
    pi q = o[n].s.top();
    o[n].s.pop();
    if (q.F != s) {
      sum += q.S;
      ppi x = search(n, q.F);
      if (x.S.F + q.S > mx) {
        sub = mx;
        mx = x.S.F + q.S;
      } else if (x.S.F + q.S > sub) {
        sub = x.S.F + q.S;
      }
      dis = max(dis, x.S.S);
      sum += x.F;
    }
  }
  dis = max(dis, mx + sub);
  p = MK(sum, MK(mx, dis));
  // cout<<p.F<<" "<<p.S.F<<" "<<p.S.S<<endl;
  return p;
}

int main() {
  while (1) {

    int n;
    cin >> n;
    if (n == 0) {
      break;
    }
    int b[1000] = {};
    int count[1000] = {};
    int sum = 0;

    for (int i = 2; i <= n; i++) {
      cin >> b[i];
      count[i]++;
      count[b[i]]++;
    }
    for (int i = 2; i <= n; i++) {
      int d;
      cin >> d;
      sum += d;
      if (count[i] != 1 && count[b[i]] != 1) {
        // cout<<i<<" "<<b[i]<<endl;
        pi p = MK(b[i], d);
        o[i].s.push(p);
        p = MK(i, d);
        o[b[i]].s.push(p);
      }
    }

    /*
    for(int i=1;i<=n;i++){
        cout<<count[i]<<" ";
    }
    cout<<endl;
     */

    ppi p;
    for (int i = 1; i <= n; i++) {
      if (count[i] != 1) {
        p = search(0, i);
        break;
      }
    }

    // cout<<p.F*2<<" "<<p.S.S<<endl;
    cout << sum + p.F * 2 - p.S.S << endl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cmath>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

int main() {
  while (1) {

    int n;
    cin >> n;
    if (n == 0) {
      break;
    }
    int a[1000][1000] = {};
    int l[1000][1000] = {};
    int b[1000] = {};
    int don = 0;
    bool done[1000] = {};
    int count[1000] = {};
    for (int i = 2; i <= n; i++) {
      cin >> b[i];
      count[b[i]]++;
      count[i]++;
    }
    int sum = 0;
    for (int i = 2; i <= n; i++) {
      int d;
      cin >> d;
      sum += d;
      a[b[i]][i] = d;
      a[i][b[i]] = d;
      l[b[i]][i] = d;
      l[i][b[i]] = d;
    }
    stack<int> o;
    for (int i = 1; i <= n; i++) {
      if (count[i] == 1) {
        done[i] = true;
        don++;
        for (int t = 1; t <= n; t++) {
          if (a[i][t] != 0) {
            o.push(i);
            o.push(t);
            a[i][t] = 0;
            a[t][i] = 0;
            break;
          }
        }
      }
    }

    while (!o.empty()) {
      int p = o.top();
      o.pop();
      count[p]--;
    }
    /*
    for(int i=1;i<=n;i++){
        cout<<count[i]<<" ";
    }
    cout<<endl;
    */
    int dis = 0;
    while (don < n - 1) {
      for (int i = 1; i <= n; i++) {
        if (count[i] == 1) {
          don++;
          for (int t = 1; t <= n; t++) {
            if (a[i][t] != 0) {
              dis += a[i][t];
              count[t]--;
              count[i]--;
              a[i][t] = 0;
              a[t][i] = 0;
              break;
            }
          }
        }
      }
    }
    dis *= 2;

    int mx = 0;
    for (int i = 1; i <= n; i++) {
      for (int t = 1; t <= n; t++) {
        for (int k = 1; k <= n; k++) {
          if (l[t][i] != 0 && l[i][k] != 0) {
            if (l[t][k] != 0) {
              l[t][k] = min(l[t][i] + l[i][k], l[t][k]);
            } else {
              l[t][k] = l[t][i] + l[i][k];
            }
            // cout<<t<<" "<<k<<" "<<l[t][k]<<endl;
          }
        }
      }
    }
    /*
    for(int i=1;i<=n;i++){
        for(int t=1;t<=n;t++){
            cout<<l[i][t]<<" ";
        }
        cout<<endl;
    }
    cout<<endl;

    for(int i=1;i<=n;i++){
        cout<<done[i]<<" ";
    }
    cout<<endl;
     */
    for (int i = 1; i <= n; i++) {
      for (int t = 1; t <= n; t++) {
        if (!done[i] && !done[t] && i != t) {
          // cout<<i<<" "<<t<<" "<<a[i][t]<<endl;
          mx = max(mx, l[i][t]);
        }
      }
    }

    cout << dis - mx + sum << endl;
  }

  return 0;
}