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

#include <iostream>
#include <set>
#include <stack>
#include <vector>

#define INF 0x3fffffff

using namespace std;

int main() {
  int n;
  while (cin >> n, n) {
    int to[900], ds[900], to_num[900] = {0}, ds_sum = 0, ans = 0;
    int w[900][900];
    set<int> sg;
    vector<int> tosg[900], dssg[900];
    fill(w[0], w[0] + 900 * 900, INF);
    for (int i = 1; i < n; i++) {
      cin >> to[i];
      to[i]--;
      to_num[to[i]]++;
      to_num[i]++;
    }
    for (int i = 1; i < n; i++) {
      cin >> ds[i];
      ds_sum += ds[i];
      if (to[i] == 0) {
        to[0] = i;
        ds[0] = ds[i];
      }
      w[i][to[i]] = ds[i];
      w[to[i]][i] = ds[i];
    }

    int ds_ha = 0;
    for (int i = 0; i < n; i++) {
      if (to_num[i] == 1) {
        ds_ha += ds[i];
        sg.insert(to[i]);
      } else {
        tosg[i].push_back(to[i]);
        dssg[i].push_back(ds[i]);
        tosg[to[i]].push_back(i);
        dssg[to[i]].push_back(ds[i]);
      }
    }

    /*		for(int i=0; i<n; i++){
                            if( !sg.count(i) ) break;
                            stack< pair<int,int> > st;
                            st.push( make_pair(i, 0) );
                            st.push( make_pair(i, 0) );
                            while( st.size() > 1 ){
                                    pair<int,int> p = st.top(); st.pop();
                                    pair<int,int> bp = st.top();
                                    for(int k=0; k<tosg[p.first].size(); k++){
                                            int j = tosg[p.first][k];
                                            if( j != bp.first ){
                                                    int sum = p.second +
       dssg[p.first][k]; if( sg.count( j ) ){ w[i][j] = min( w[i][j], sum);
                                                            w[j][i] = min(
       w[j][i], sum);
                                                    }
                                                    st.push(
       make_pair(tosg[p.first][k], sum) );
                                            }
                                    }
                            }
                    }*/

    for (int i = 1; i < n; i++) {
      int j = to[i], sum = ds[i];
      if (sg.count(i))
        while (1) {
          //			if( sg.count(j) || j == 0 ){
          w[i][j] = min(w[i][j], sum);
          w[j][i] = min(w[j][i], sum);
          for (int k = 0; k < tosg[j].size(); k++) {
            w[i][tosg[j][k]] = min(w[i][tosg[j][k]], sum + dssg[j][k]);
            w[tosg[j][k]][i] = min(w[i][tosg[j][k]], sum + dssg[j][k]);
          }
          tosg[j].push_back(i);
          dssg[j].push_back(sum);
          //			}
          if (j == 0) {
            break;
          } else {
            sum += ds[j];
            j = to[j];
          }
        }
    }
    // for(int k=0; k<n; k++){
    //	for(int i=0; i<n; i++){
    //		for(int j=0; j<n; j++){
    //			w[i][j] = min(w[i][j], w[i][k] + w[k][j]);
    //		}
    //	}
    // }
    int maxlen = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < n; j++) {
        if (sg.count(i) && sg.count(j) && i != j && w[i][j] < INF) {
          maxlen = max(maxlen, w[i][j]);
        }
      }
    }
    ans = (3 * ds_sum) - (2 * ds_ha) - maxlen;
    //	cout << ds_sum << " " << ds_ha << " " << maxlen << endl;
    cout << ans << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <iostream>
#include <set>
#include <vector>

#define INF 0x3fffffff

using namespace std;

int main() {
  int n;
  while (cin >> n, n) {
    int to[900], ds[900], to_num[900] = {0}, ds_sum = 0, ans = 0;
    int w[900][900];
    set<int> sg;
    fill(w[0], w[0] + 900 * 900, INF);
    for (int i = 1; i < n; i++) {
      cin >> to[i];
      to[i]--;
      to_num[to[i]]++;
      to_num[i]++;
    }
    for (int i = 1; i < n; i++) {
      cin >> ds[i];
      ds_sum += ds[i];
      if (to[i] == 0) {
        to[0] = i;
        ds[0] = ds[i];
      }
      w[i][to[i]] = ds[i];
      w[to[i]][i] = ds[i];
    }
    int ds_ha = 0;
    for (int i = 0; i < n; i++) {
      if (to_num[i] == 1) {
        ds_ha += ds[i];
        sg.insert(to[i]);
      }
    }
    for (int k = 0; k < n; k++) {
      for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
          w[i][j] = min(w[i][j], w[i][k] + w[k][j]);
        }
      }
    }
    int maxlen = 0;
    for (int i = 0; i < n; i++) {
      for (int j = 0; j < n; j++) {
        if (sg.count(i) && sg.count(j) && i != j && w[i][j] < INF) {
          maxlen = max(maxlen, w[i][j]);
        }
      }
    }
    ans = (3 * ds_sum) - (2 * ds_ha) - maxlen;
    //	cout << ds_sum << " " << ds_ha << " " << maxlen << endl;
    cout << ans << endl;
  }
  return 0;
}