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

#include <bits/stdc++.h>

using namespace std;

#define int long long

#define reps(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define rep(i, n) reps(i, 0, n)

#define MAX_N 500005

using vint = vector<int>;
using Graph = vector<vint>;

int n;
Graph graph;

int dp1[MAX_N];
void dfs1(int root) {
  memset(dp1, -1, sizeof(dp1));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp1[u] == -1) {
      dp1[u] = 2 - graph[u].size();
      for (int v : graph[u]) {
        stk.push(v);
      }
    } else {
      for (int v : graph[u]) {
        dp1[u] += dp1[v];
      }
      stk.pop();
    }
  }
}

int dp2[MAX_N];
int depth[MAX_N];
void dfs2(int root) {
  memset(dp2, -1, sizeof(dp2));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp2[u] == -1) {
      dp2[u] = dp1[u];
      for (int v : graph[u]) {
        depth[v] = depth[u] + 1;
        stk.push(v);
      }
    } else {
      for (int v : graph[u]) {
        dp2[u] += dp2[v];
      }
      stk.pop();
    }
  }
}

int dp3[MAX_N];
void dfs3(int root) {
  memset(dp3, -1, sizeof(dp3));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp3[u] == -1) {
      dp3[u] = 0;
      for (int v : graph[u]) {
        stk.push(v);
      }
    } else {
      if (graph[u].size() == 2) {
        int l = graph[u][0], r = graph[u][1];
        dp3[u] =
            min({dp3[l] + 1 + dp3[r] + 1 + depth[u],
                 dp3[l] + 1 + dp2[r] + dp1[r], dp2[l] + dp1[l] + dp3[r] + 1});
      } else if (graph[u].size() == 1) {
        int ch = graph[u][0];
        dp3[u] = min({dp3[ch] + 1 + 1 + depth[u], dp3[ch] + 1 + 1,
                      1 + dp2[ch] + dp1[ch]});
      } else {
        dp3[u] = dp1[u];
      }
      stk.pop();
    }
  }
}

signed main() {
  cin >> n;
  graph.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    if (a != n - 1)
      graph[i].push_back(a);
    if (b != n - 1)
      graph[i].push_back(b);
  }
  dfs1(0);
  dfs2(0);
  dfs3(0);
  cout << dp3[0] << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define int long long

#define reps(i, m, n) for (int i = (int)(m); i < (int)(n); ++i)
#define rep(i, n) reps(i, 0, n)

#define MAX_N 500005

using vint = vector<int>;
using Graph = vector<vint>;

int n;
Graph graph;

int dp1[MAX_N];
void dfs1(int root) {
  memset(dp1, -1, sizeof(dp1));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp1[u] == -1) {
      dp1[u] = 2 - graph[u].size();
      for (int v : graph[u]) {
        stk.push(v);
      }
    } else {
      for (int v : graph[u]) {
        dp1[u] += dp1[v];
      }
      stk.pop();
    }
  }
}

int dp2[MAX_N];
int depth[MAX_N];
void dfs2(int root) {
  memset(dp2, -1, sizeof(dp2));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp2[u] == -1) {
      dp2[u] = dp1[u];
      for (int v : graph[u]) {
        depth[v] = depth[u] + 1;
        stk.push(v);
      }
    } else {
      for (int v : graph[u]) {
        dp2[u] += dp2[v];
      }
      stk.pop();
    }
  }
}

int dp3[MAX_N];
void dfs3(int root) {
  memset(dp3, -1, sizeof(dp3));
  stack<int> stk;
  stk.push(root);
  while (!stk.empty()) {
    int u = stk.top();
    if (dp3[u] == -1) {
      dp3[u] = 0;
      for (int v : graph[u]) {
        stk.push(v);
      }
    } else {
      if (graph[u].size() == 2) {
        int l = graph[u][0], r = graph[u][1];
        dp3[u] =
            min({dp3[l] + 1 + dp3[r] + 1 + depth[u],
                 dp3[l] + 1 + dp2[r] + dp1[r], dp2[l] + dp1[l] + dp3[r] + 1});
      } else if (graph[u].size() == 1) {
        int ch = graph[u][0];
        dp3[u] = min({dp3[ch] + 1 + 1 + depth[u], dp3[ch] + 1 + 1,
                      1 + dp2[ch] + dp1[ch]});
      } else {
        dp3[u] = dp1[u];
      }
    }
  }
}

signed main() {
  cin >> n;
  graph.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    --a, --b;
    if (a != n - 1)
      graph[i].push_back(a);
    if (b != n - 1)
      graph[i].push_back(b);
  }
  dfs1(0);
  dfs2(0);
  dfs3(0);
  cout << dp3[0] << endl;

  return 0;
}