Problem p03204 (C++)

Buggy submission (status: Time Limit Exceeded):

#include <bits/stdc++.h>

#define ln '\n'
#define all(dat) dat.begin(), dat.end()
#define loop(i, to) for (int i = 0; i < to; ++i)
#define cont(i, to) for (int i = 1; i <= to; ++i)
#define circ(i, fm, to) for (int i = fm; i <= to; ++i)
#define foreach(i, dat)                                                        \
  for (__typeof(dat.begin()) i = dat.begin(); i != dat.end(); ++i)

typedef long long num;

using namespace std;

const int nsz = 2e5;
int n, rt = 1, mx[nsz + 5], pr[nsz + 5], dp[nsz + 5][3], ans[nsz + 5];
vector<int> g[nsz + 5];

int inline qry0(int u, int w) {
  int cnt = 1;
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u] || v > w)
      continue;
    cnt += qry0(v, w);
  }
  return cnt;
}

int inline qry1(int u, int w) {
  if (u > w)
    return 0;
  int cnt = 1;
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    cnt += qry1(v, w);
  }
  return cnt;
}

void dfs(int u = rt) {
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    pr[v] = u;
    mx[v] = max(mx[u], u);
    dfs(v);
  }
  if (u > mx[u]) {
    dp[u][0] = 1;
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][1];
      }
    }
  } else if (u < mx[u] && pr[u] < mx[u]) {
    dp[u][0] = dp[u][1] = dp[u][2] = 1;
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][0];
      }
      if (v <= mx[pr[u]]) {
        dp[u][1] += dp[v][1];
      }
      if (v <= mx[mx[pr[u]]]) {
        dp[u][2] += dp[v][2];
      }
    }
  } else {
    dp[u][0] = 1;
    dp[u][1] = u <= mx[pr[u]];
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][0];
      }
      if (u <= mx[pr[u]] && v <= mx[pr[u]]) {
        dp[u][1] += dp[v][2];
      }
    }
    // dp[u][1] = qry1(u, mx[pr[u]]);
  }
  if (u == rt)
    return;
  if (u > mx[u]) {
    ans[u] = dp[u][0];
    // ans[u] = qry0(u, mx[u]);
  } else if (pr[u] == mx[u]) {
    ans[u] = dp[u][0] - dp[u][1];
    // ans[u] = qry0(u, mx[u]) - qry1(u, mx[pr[u]]);
  } else {
    ans[u] = 0;
  }
}

void pre(int u = rt) {
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    ans[v] += ans[u];
    pre(v);
  }
}

int main() {
  // freopen("dat.in", "r", stdin);
  // freopen("dat.out", "w", stdout);
  scanf("%d", &n);
  cont(i, n - 1) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs();
  cont(u, n) {
    if (dp[u][0] != qry0(u, mx[u]) || dp[u][1] != qry1(u, mx[pr[u]])) {
      cerr << u << ": " << dp[u][0] << ' ' << qry0(u, mx[u]) << ' ' << dp[u][1]
           << ' ' << qry1(u, mx[pr[u]]) << ln;
    }
  }
  pre();
  circ(u, 2, n) { cout << ans[u] << ' '; }
  cout << ln;
}

Fixed submission:

#include <bits/stdc++.h>

#define ln '\n'
#define all(dat) dat.begin(), dat.end()
#define loop(i, to) for (int i = 0; i < to; ++i)
#define cont(i, to) for (int i = 1; i <= to; ++i)
#define circ(i, fm, to) for (int i = fm; i <= to; ++i)
#define foreach(i, dat)                                                        \
  for (__typeof(dat.begin()) i = dat.begin(); i != dat.end(); ++i)

typedef long long num;

using namespace std;

const int nsz = 2e5;
int n, rt = 1, mx[nsz + 5], pr[nsz + 5], dp[nsz + 5][3], ans[nsz + 5];
vector<int> g[nsz + 5];

int inline qry0(int u, int w) {
  int cnt = 1;
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u] || v > w)
      continue;
    cnt += qry0(v, w);
  }
  return cnt;
}

int inline qry1(int u, int w) {
  if (u > w)
    return 0;
  int cnt = 1;
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    cnt += qry1(v, w);
  }
  return cnt;
}

void dfs(int u = rt) {
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    pr[v] = u;
    mx[v] = max(mx[u], u);
    dfs(v);
  }
  if (u > mx[u]) {
    dp[u][0] = 1;
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][1];
      }
    }
  } else if (u < mx[u] && pr[u] < mx[u]) {
    dp[u][0] = dp[u][1] = dp[u][2] = 1;
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][0];
      }
      if (v <= mx[pr[u]]) {
        dp[u][1] += dp[v][1];
      }
      if (v <= mx[mx[pr[u]]]) {
        dp[u][2] += dp[v][2];
      }
    }
  } else {
    dp[u][0] = 1;
    dp[u][1] = u <= mx[pr[u]];
    loop(i, g[u].size()) {
      int v = g[u][i];
      if (v == pr[u])
        continue;
      if (v <= mx[u]) {
        dp[u][0] += dp[v][0];
      }
      if (u <= mx[pr[u]] && v <= mx[pr[u]]) {
        dp[u][1] += dp[v][2];
      }
    }
    // dp[u][1] = qry1(u, mx[pr[u]]);
  }
  if (u == rt)
    return;
  if (u > mx[u]) {
    ans[u] = dp[u][0];
    // ans[u] = qry0(u, mx[u]);
  } else if (pr[u] == mx[u]) {
    ans[u] = dp[u][0] - dp[u][1];
    // ans[u] = qry0(u, mx[u]) - qry1(u, mx[pr[u]]);
  } else {
    ans[u] = 0;
  }
}

void pre(int u = rt) {
  loop(i, g[u].size()) {
    int v = g[u][i];
    if (v == pr[u])
      continue;
    ans[v] += ans[u];
    pre(v);
  }
}

int main() {
  // freopen("dat.in", "r", stdin);
  // freopen("dat.out", "w", stdout);
  scanf("%d", &n);
  cont(i, n - 1) {
    int u, v;
    scanf("%d%d", &u, &v);
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs();
  // cont (u, n) {
  //     if (dp[u][0] != qry0(u, mx[u]) || dp[u][1] != qry1(u, mx[pr[u]])) {
  //         cerr << u << ": " << dp[u][0] << ' ' << qry0(u, mx[u]) << ' ' <<
  //         dp[u][1] << ' ' << qry1(u, mx[pr[u]]) << ln;
  //     }
  // }
  pre();
  circ(u, 2, n) { cout << ans[u] << ' '; }
  cout << ln;
}