Problem p02710 (C++)

Buggy submission (status: Runtime Error):

#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 = int(2e5);
bool vis[nsz + 5];
int n, rt = 1, a[nsz + 5], pr[nsz + 5];
vector<int> p[nsz + 5];

struct link_cut_tree {
#define dir(u) (rs[pr[u]] == u)
#define s(u, d) (d ? rs[u] : ls[u])
#define is_rt(u) (ls[pr[u]] != u && rs[pr[u]] != u)

  bool rev[nsz + 5];
  int sz, pr[nsz + 5], ls[nsz + 5], rs[nsz + 5];
  int siz[nsz + 5], vsz[nsz + 5];

  void inline up(int u) { siz[u] = siz[ls[u]] + siz[rs[u]] + 1 + vsz[u]; }

  void inline down(int u) {
    if (rev[u]) {
      int v1 = ls[u], v2 = rs[u];
      if (v1)
        swap(ls[v1], rs[v1]), rev[v1] ^= 1;
      if (v2)
        swap(ls[v2], rs[v2]), rev[v2] ^= 1;
      rev[u] = 0;
    }
  }

  void inline push_down(int u) {
    static int stk[nsz + 5];
    int s = 0;
    for (; !is_rt(u); stk[s++] = u, u = pr[u])
      ;
    for (stk[s++] = u; s; down(stk[--s]))
      ;
  }

  void inline rotate(int u) {
    int p = pr[u], g = pr[p], d = dir(u), v = s(u, !d);
    !is_rt(p) && (s(g, dir(p)) = u);
    pr[u] = g;
    s(p, d) = v, pr[v] = p;
    s(u, !d) = p, pr[p] = u;
    up(p), up(u);
  }

  void inline splay(int u) {
    push_down(u);
    for (int p = pr[u]; !is_rt(u); rotate(u), p = pr[u]) {
      if (!is_rt(p))
        rotate(dir(u) != dir(p) ? u : p);
    }
  }

  int inline access(int u) {
    int v = 0;
    for (; u; v = u, u = pr[u]) {
      splay(u);
      vsz[u] += siz[rs[u]] - siz[v];
      rs[u] = v, up(u);
    }
    return v;
  }

  void inline set_rt(int u) {
    u = access(u);
    swap(ls[u], rs[u]), rev[u] ^= 1;
  }

  int inline rt_of(int u) {
    u = access(u);
    for (; ls[u]; down(u), u = ls[u])
      ;
    splay(u);
    return u;
  }

  bool inline link(int u, int v) {
    set_rt(u);
    if (rt_of(v) == u)
      return 0;
    splay(v);
    rs[v] = u, pr[u] = v;
    up(v);
    return 1;
  }

  bool inline cut(int u, int v) {
    set_rt(u), access(v), splay(v);
    if (ls[v] != u)
      return 0;
    pr[u] = ls[v] = 0;
    up(u);
    return 1;
  }

  int inline node() {
    int u = ++sz;
    ls[u] = rs[u] = siz[u] = vsz[u] = rev[u] = 0;
    return u;
  }
};
link_cut_tree tr;

struct graph {
  int sz, hd[nsz + 5];

  struct edge {
    int nxt, to;

    edge(int nxt = 0, int to = 0) { this->nxt = nxt, this->to = to; }
  } e[2 * nsz + 5];

  inline edge &operator[](int id) { return e[id]; }

  void inline link(int u, int v) {
    e[++sz] = edge(hd[u], v), hd[u] = sz;
    e[++sz] = edge(hd[v], u), hd[v] = sz;
  }
};
graph g;

int main() {
  scanf("%d", &n);
  cont(i, n) {
    scanf("%d", &a[i]);
    tr.node();
    p[a[i]].push_back(i);
  }
  cont(i, n - 1) {
    int u, v;
    scanf("%d%d", &u, &v);
    g.link(u, v);
    tr.link(u, v);
  }

  cont(c, n) {
    if (!p[c].size()) {
      printf("0\n");
      continue;
    }
    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to;
        tr.cut(u, v);
      }
    }

    num res = (num)n * (n + 1) / 2;
    static int stk[nsz + 5];
    int s = 0;
    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to, p = tr.rt_of(v), sz = tr.siz[p];
        if (v == pr[u] || a[v] == c || vis[p])
          continue;
        res -= (num)sz * (sz + 1) / 2;
        vis[p] = 1;
        stk[s++] = p;
      }
    }
    printf("%lld\n", res);

    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to;
        tr.link(u, v);
      }
    }
    for (; s; vis[stk[--s]] = 0)
      ;
  }
}

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 = int(2e5);
bool vis[nsz + 5];
int n, rt = 1, a[nsz + 5], pr[nsz + 5];
vector<int> p[nsz + 5];

struct link_cut_tree {
#define dir(u) (rs[pr[u]] == u)
#define s(u, d) (d ? rs[u] : ls[u])
#define is_rt(u) (ls[pr[u]] != u && rs[pr[u]] != u)

  bool rev[nsz + 5];
  int sz, pr[nsz + 5], ls[nsz + 5], rs[nsz + 5];
  int siz[nsz + 5], vsz[nsz + 5];

  void inline up(int u) { siz[u] = siz[ls[u]] + siz[rs[u]] + 1 + vsz[u]; }

  void inline down(int u) {
    if (rev[u]) {
      int v1 = ls[u], v2 = rs[u];
      if (v1)
        swap(ls[v1], rs[v1]), rev[v1] ^= 1;
      if (v2)
        swap(ls[v2], rs[v2]), rev[v2] ^= 1;
      rev[u] = 0;
    }
  }

  void inline push_down(int u) {
    static int stk[nsz + 5];
    int s = 0;
    for (; !is_rt(u); stk[s++] = u, u = pr[u])
      ;
    for (stk[s++] = u; s; down(stk[--s]))
      ;
  }

  void inline rotate(int u) {
    int p = pr[u], g = pr[p], d = dir(u), v = s(u, !d);
    !is_rt(p) && (s(g, dir(p)) = u);
    pr[u] = g;
    s(p, d) = v, pr[v] = p;
    s(u, !d) = p, pr[p] = u;
    up(p), up(u);
  }

  void inline splay(int u) {
    push_down(u);
    for (int p = pr[u]; !is_rt(u); rotate(u), p = pr[u]) {
      if (!is_rt(p))
        rotate(dir(u) != dir(p) ? u : p);
    }
  }

  int inline access(int u) {
    int v = 0;
    for (; u; v = u, u = pr[u]) {
      splay(u);
      vsz[u] += siz[rs[u]] - siz[v];
      rs[u] = v, up(u);
    }
    return v;
  }

  void inline set_rt(int u) {
    access(u), splay(u);
    swap(ls[u], rs[u]), rev[u] ^= 1;
  }

  int inline rt_of(int u) {
    u = access(u);
    for (; ls[u]; down(u), u = ls[u])
      ;
    splay(u);
    return u;
  }

  bool inline link(int u, int v) {
    set_rt(u);
    if (rt_of(v) == u)
      return 0;
    splay(v);
    rs[v] = u, pr[u] = v;
    up(v);
    return 1;
  }

  bool inline cut(int u, int v) {
    set_rt(u), access(v), splay(v);
    if (ls[v] != u)
      return 0;
    pr[u] = ls[v] = 0;
    up(u);
    return 1;
  }

  int inline node() {
    int u = ++sz;
    ls[u] = rs[u] = siz[u] = vsz[u] = rev[u] = 0;
    return u;
  }
};
link_cut_tree tr;

struct graph {
  int sz, hd[nsz + 5];

  struct edge {
    int nxt, to;

    edge(int nxt = 0, int to = 0) { this->nxt = nxt, this->to = to; }
  } e[2 * nsz + 5];

  inline edge &operator[](int id) { return e[id]; }

  void inline link(int u, int v) {
    e[++sz] = edge(hd[u], v), hd[u] = sz;
    e[++sz] = edge(hd[v], u), hd[v] = sz;
  }
};
graph g;

int main() {
  scanf("%d", &n);
  cont(i, n) {
    scanf("%d", &a[i]);
    tr.node();
    p[a[i]].push_back(i);
  }
  cont(i, n - 1) {
    int u, v;
    scanf("%d%d", &u, &v);
    g.link(u, v);
    tr.link(u, v);
  }

  cont(c, n) {
    if (!p[c].size()) {
      printf("0\n");
      continue;
    }
    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to;
        tr.cut(u, v);
      }
    }

    num res = (num)n * (n + 1) / 2;
    static int stk[nsz + 5];
    int s = 0;
    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to, p = tr.rt_of(v), sz = tr.siz[p];
        if (v == pr[u] || a[v] == c || vis[p])
          continue;
        res -= (num)sz * (sz + 1) / 2;
        vis[p] = 1;
        stk[s++] = p;
      }
    }
    printf("%lld\n", res);

    loop(i, p[c].size()) {
      int u = p[c][i];
      for (int i = g.hd[u]; i; i = g[i].nxt) {
        int v = g[i].to;
        tr.link(u, v);
      }
    }
    for (; s; vis[stk[--s]] = 0)
      ;
  }
}