Problem p02710 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

using namespace std;

int n;
int c[200010];
vector<int> g[200010];
int dep[200010];
int fa[200010][20];
int sz[200010];
long long ans[200010];
vector<int> stck[200010];

inline void dfs(int x) {
  for (int i = 1; i < 20; i++)
    fa[x][i] = fa[fa[x][i - 1]][i - 1];

  sz[x] = 1;
  for (auto ch : g[x])
    if (ch != fa[x][0]) {
      dep[ch] = dep[x] + 1;
      fa[ch][0] = x;
      dfs(ch);
      sz[x] += sz[ch];
    }
}

inline void swim(int &x, int h) {
  for (int i = 0; h > 0; i++) {
    if (h & 1)
      x = fa[x][i];
    h >>= 1;
  }
}

int N;
inline void dfs1(int x) {
  for (auto ch : g[x]) {
    if (ch != fa[x][0]) {
      dfs1(ch);

      int num = sz[ch];
      while (!stck[c[x]].empty()) {
        int temp = *stck[c[x]].rbegin();
        swim(temp, dep[temp] - dep[x]);
        if (temp == x) {
          temp = *stck[c[x]].rbegin();
          num -= sz[temp];
          stck[c[x]].pop_back();
        } else {
          break;
        }
      }
      ans[c[x]] -= (long long)num * (num + 1) / 2LL;
    }
  }
  stck[c[x]].push_back(x);
}

int main() {
  // freopen("f.in","r",stdin);
  scanf("%d", &n);
  for (int i = 1; i <= n; i++) {
    scanf("%d", c + i);
  }
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    g[x].push_back(y);
    g[y].push_back(x);
  }
  fa[1][0] = 1;
  dep[1] = 1;
  dfs(1);

  for (int i = 1; i <= n; i++) {
    ans[i] = (long long)n * (n + 1) / 2LL;
  }
  dfs1(1);

  for (int i = 1; i <= n; i++) {
    if (i == c[1])
      continue;

    int num = n;
    while (!stck[i].empty()) {
      int temp = *stck[i].rbegin();
      num -= sz[temp];
      stck[i].pop_back();
    }

    ans[i] -= (long long)num * (num + 1) / 2LL;
  }

  for (int i = 1; i <= n; i++) {
    printf("%lld\\n", ans[i]);
  }

  return 0;
}

Error: 0

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

int n;
int c[200010];
vector<int> g[200010];
int dep[200010];
int fa[200010][20];
int sz[200010];
long long ans[200010];
vector<int> stck[200010];

inline void dfs(int x) {
  for (int i = 1; i < 20; i++)
    fa[x][i] = fa[fa[x][i - 1]][i - 1];

  sz[x] = 1;
  for (auto ch : g[x])
    if (ch != fa[x][0]) {
      dep[ch] = dep[x] + 1;
      fa[ch][0] = x;
      dfs(ch);
      sz[x] += sz[ch];
    }
}

inline void swim(int &x, int h) {
  for (int i = 0; h; i++) {
    if (h & 1)
      x = fa[x][i];
    h >>= 1;
  }
}

inline void dfs1(int x) {
  for (auto ch : g[x])
    if (ch != fa[x][0]) {
      dfs1(ch);

      int num = sz[ch];
      while (!stck[c[x]].empty()) {
        int temp = *stck[c[x]].rbegin();
        swim(temp, dep[temp] - dep[x]);
        if (temp == x) {
          temp = *stck[c[x]].rbegin();
          num -= sz[temp];
          stck[c[x]].pop_back();
        } else {
          break;
        }
      }
      ans[c[x]] -= (long long)num * (num + 1) / 2LL;
    }
  stck[c[x]].push_back(x);
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i <= n; i++) {
    scanf("%d", c + i);
  }
  for (int i = 1; i < n; i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    g[x].push_back(y);
    g[y].push_back(x);
  }
  fa[1][0] = 1;
  dep[1] = 1;
  dfs(1);

  for (int i = 1; i <= n; i++) {
    ans[i] = (long long)n * (n + 1) / 2LL;
  }

  dfs1(1);

  // mdf

  for (int i = 1; i <= n; i++) {
    if (i == c[1])
      continue;

    int num = n;
    while (!stck[i].empty()) {
      int temp = *stck[i].rbegin();
      num -= sz[temp];
      stck[i].pop_back();
    }

    ans[i] -= (long long)num * (num + 1) / 2LL;
  }

  for (int i = 1; i <= n; i++) {
    printf("%lld\\n", ans[i]);
  }

  return 0;
}