Language: C

#pragma GCC optimize ("Ofast")
#include<bits/stdc++.h>
using namespace std;
void *wmem;
char memarr[96000000];
template<class T> inline void walloc1d(T **arr, int x, void **mem = &wmem){
  static int skip[16] = {0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  (*mem) = (void*)( ((char*)(*mem)) + skip[((unsigned long long)(*mem)) & 15] );
  (*arr)=(T*)(*mem);
  (*mem)=((*arr)+x);
}
inline int my_getchar_unlocked(){
  static char buf[1048576];
  static int s = 1048576;
  static int e = 1048576;
  if(s == e && e == 1048576){
    e = fread_unlocked(buf, 1, 1048576, stdin);
    s = 0;
  }
  if(s == e){
    return EOF;
  }
  return buf[s++];
}
inline void rd(int &x){
  int k;
  int m=0;
  x=0;
  for(;;){
    k = my_getchar_unlocked();
    if(k=='-'){
      m=1;
      break;
    }
    if('0'<=k&&k<='9'){
      x=k-'0';
      break;
    }
  }
  for(;;){
    k = my_getchar_unlocked();
    if(k<'0'||k>'9'){
      break;
    }
    x=x*10+k-'0';
  }
  if(m){
    x=-x;
  }
}
struct MY_WRITER{
  char buf[1048576];
  int s;
  int e;
  MY_WRITER(){
    s = 0;
    e = 1048576;
  }
  ~MY_WRITER(){
    if(s){
      fwrite_unlocked(buf, 1, s, stdout);
    }
  }
}
;
MY_WRITER MY_WRITER_VAR;
void my_putchar_unlocked(int a){
  if(MY_WRITER_VAR.s == MY_WRITER_VAR.e){
    fwrite_unlocked(MY_WRITER_VAR.buf, 1, MY_WRITER_VAR.s, stdout);
    MY_WRITER_VAR.s = 0;
  }
  MY_WRITER_VAR.buf[MY_WRITER_VAR.s++] = a;
}
inline void wt_L(char a){
  my_putchar_unlocked(a);
}
inline void wt_L(long long x){
  int s=0;
  int m=0;
  char f[20];
  if(x<0){
    m=1;
    x=-x;
  }
  while(x){
    f[s++]=x%10;
    x/=10;
  }
  if(!s){
    f[s++]=0;
  }
  if(m){
    my_putchar_unlocked('-');
  }
  while(s--){
    my_putchar_unlocked(f[s]+'0');
  }
}
struct graph{
  int N;
  int *es;
  int **edge;
  void setEdge(int N__, int M, int A[], int B[], void **mem = &wmem){
    int i;
    N = N__;
    walloc1d(&es, N, mem);
    walloc1d(&edge, N, mem);
    for(i=(0);i<(N);i++){
      es[i] = 0;
    }
    for(i=(0);i<(M);i++){
      es[A[i]]++;
      es[B[i]]++;
    }
    for(i=(0);i<(N);i++){
      walloc1d(&edge[i], es[i], mem);
    }
    for(i=(0);i<(N);i++){
      es[i] = 0;
    }
    for(i=(0);i<(M);i++){
      edge[A[i]][es[A[i]]++] = B[i];
      edge[B[i]][es[B[i]]++] = A[i];
    }
  }
  void SubTreeSize(int root, int res[], void *mem = wmem){
    int i;
    int j;
    int k;
    int m;
    int *q;
    int qs = 0;
    int qe = 1;
    walloc1d(&q,N,&mem);
    for(i=(0);i<(N);i++){
      res[i] = -1;
    }
    res[root] = 0;
    q[0] = root;
    while(qs < qe){
      i = q[qs++];
      for(j=(0);j<(es[i]);j++){
        k = edge[i][j];
        if(res[k]==0){
          continue;
        }
        res[k] = 0;
        q[qe++] = k;
      }
    }
    for(m=(N)-1;m>=(0);m--){
      i = q[m];
      res[i] = 1;
      for(j=(0);j<(es[i]);j++){
        k = edge[i][j];
        res[i] += res[k];
      }
    }
  }
  int preorder(int res[], int root = 0, void *mem=wmem){
    int i;
    int j;
    int k;
    int sts;
    int sz = 0;
    long long *st;
    char *vis;
    walloc1d(&vis, N, &mem);
    walloc1d(&st, N, &mem);
    sts = 0;
    st[sts++] = ((long long)root) << 32;
    for(i=(0);i<(N);i++){
      vis[i] = 0;
    }
    vis[root] = 1;
    while(sts){
      i = st[--sts] >> 32;
      j = st[sts] & 2147483647;
      if(j==0){
        res[sz++] = i;
      }
      while(j < es[i]){
        k = edge[i][j++];
        if(vis[k]){
          continue;
        }
        if(j < es[i]){
          st[sts++] = (((long long)i) << 32) + j;
        }
        vis[k] = 1;
        st[sts++] = ((long long)k) << 32;
        break;
      }
    }
    return sz;
  }
}
;
struct HLD{
  int N;
  int *es;
  int **edge;
  int *group;
  int *groupind;
  int groupNum;
  int *groupSize;
  int **groupNode;
  int *groupUpNode;
  int *groupDepth;
  void init(graph g, void **mem = &wmem){
    init(g.N, g.es, g.edge, mem);
  }
  void init(int N__, int *es__, int **edge__, void **mem = &wmem){
    int i;
    int j;
    int k;
    int x;
    int y;
    int mx;
    int *q;
    int q_st;
    int q_ed;
    int *sz;
    char *vis;
    void *tmpmem;
    N = N__;
    es = es__;
    edge = edge__;
    walloc1d(&group, N, mem);
    walloc1d(&groupind, N, mem);
    tmpmem = *mem;
    walloc1d(&q, N, &tmpmem);
    walloc1d(&sz, N, &tmpmem);
    walloc1d(&vis, N, &tmpmem);
    for(i=(0);i<(N);i++){
      vis[i] = 0;
    }
    q_st = 0;
    q_ed = 1;
    q[0] = 0;
    vis[0] = 1;
    while(q_st < q_ed){
      i = q[q_st++];
      for(j=(0);j<(es[i]);j++){
        k = edge[i][j];
        if(!vis[k]){
          vis[k] = 1;
          q[q_ed++] = k;
        }
      }
    }
    for(i=(0);i<(N);i++){
      sz[i] = 0;
    }
    for(j=N-1;j>=0;j--){
      i = q[j];
      sz[i] = 1;
      for(k=(0);k<(es[i]);k++){
        sz[i] += sz[edge[i][k]];
      }
    }
    for(i=(0);i<(N);i++){
      group[i] = -1;
    }
    groupNum = 0;
    for(j=(0);j<(N);j++){
      i = q[j];
      if(group[i]>=0){
        continue;
      }
      group[i] = groupNum++;
      groupind[i] = 0;
      for(;;){
        mx = -1;
        for(k=(0);k<(es[i]);k++){
          if(group[edge[i][k]] != -1){
            continue;
          }
          if(mx==-1){
            mx = k;
          }
          else if(sz[edge[i][k]] > sz[edge[i][mx]]){
            mx = k;
          }
        }
        if(mx==-1){
          break;
        }
        group[edge[i][mx]] = group[i];
        groupind[edge[i][mx]] = groupind[i]+1;
        i = edge[i][mx];
      }
    }
    walloc1d(&groupSize, groupNum, mem);
    walloc1d(&groupUpNode, groupNum, mem);
    walloc1d(&groupDepth, groupNum, mem);
    for(i=(0);i<(groupNum);i++){
      groupSize[i] = 0;
    }
    for(i=(0);i<(N);i++){
      groupSize[group[i]]++;
    }
    walloc1d(&groupNode, groupNum, mem);
    for(i=(0);i<(groupNum);i++){
      walloc1d(&groupNode[i], groupSize[i], mem);
    }
    for(i=(0);i<(N);i++){
      groupNode[group[i]][groupind[i]] = i;
    }
    for(i=(0);i<(groupNum);i++){
      groupDepth[i] = -1;
    }
    groupUpNode[0] = -1;
    groupDepth[0] = 0;
    for(x=(0);x<(groupNum);x++){
      for(y=(0);y<(groupSize[x]);y++){
        i = groupNode[x][y];
        for(j=(0);j<(es[i]);j++){
          k = edge[i][j];
          if(x != group[k] && groupDepth[group[k]]==-1){
            groupUpNode[group[k]] = i;
            groupDepth[group[k]] = groupDepth[x] + 1;
          }
        }
      }
    }
  }
  int lca(int x, int y){
    int x1;
    int y1;
    int x2;
    int y2;
    x1 = group[x];
    x2 = groupind[x];
    y1 = group[y];
    y2 = groupind[y];
    while(groupDepth[x1] > groupDepth[y1]){
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
    }
    while(groupDepth[x1] < groupDepth[y1]){
      y = groupUpNode[y1];
      y1 = group[y];
      y2 = groupind[y];
    }
    while(x1 != y1){
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
      y = groupUpNode[y1];
      y1 = group[y];
      y2 = groupind[y];
    }
    if(x2 <= y2){
      return x;
    }
    return y;
  }
  int depth(int x){
    int x1;
    int x2;
    int res = 0;
    x1 = group[x];
    x2 = groupind[x];
    while(groupUpNode[x1] != -1){
      res += x2 + 1;
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
    }
    return res + x2;
  }
  int dist(int x, int y){
    int x1;
    int y1;
    int x2;
    int y2;
    int res = 0;
    x1 = group[x];
    x2 = groupind[x];
    y1 = group[y];
    y2 = groupind[y];
    while(groupDepth[x1] > groupDepth[y1]){
      res += x2 + 1;
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
    }
    while(groupDepth[x1] < groupDepth[y1]){
      res += y2 + 1;
      y = groupUpNode[y1];
      y1 = group[y];
      y2 = groupind[y];
    }
    while(x1 != y1){
      res += x2 + y2 + 2;
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
      y = groupUpNode[y1];
      y1 = group[y];
      y2 = groupind[y];
    }
    if(x2 <= y2){
      return res + y2 - x2;
    }
    return res + x2 - y2;
  }
  int up(int x){
    int x1 = group[x];
    int x2 = groupind[x];
    if(x2==0){
      return groupUpNode[x1];
    }
    return groupNode[x1][x2-1];
  }
  int up(int x, int d){
    int x1 = group[x];
    int x2 = groupind[x];
    while(d > x2){
      if(groupUpNode[x1]==-1){
        return -1;
      }
      d -= x2 + 1;
      x = groupUpNode[x1];
      x1 = group[x];
      x2 = groupind[x];
    }
    return groupNode[x1][x2-d];
  }
}
;
int N;
int A[200000];
int B[200000];
int C[200000];
graph g;
HLD hld;
int pre[200000];
int prec[200000];
int rev[200000];
int ps;
int sz[200000];
vector<int> b[200000];
long long res[200000];
int solve(int r, int &ind, vector<int> &b, long long &res){
  int i;
  int k = 0;
  int nx;
  long long s = 0;
  long long em;
  if(r==-1){
    while(ind < b.size()){
      i = ind++;
      s += solve(b[i], ind, b, res);
    }
    em = N - s;
    res -= em * (em+1) / 2;
  }
  else{
    for(k=(0);k<(g.es[r]);k++){
      if(g.edge[r][k] != hld.up(r)){
        nx = g.edge[r][k];
        while(ind < b.size() && hld.lca(nx, b[ind]) == nx){
          i = ind++;
          s += solve(b[i], ind, b, res);
        }
        em = sz[nx] - s;
        res -= em * (em+1) / 2;
        s = 0;
      }
    }
  }
  if(r==-1){
    return -1;
  }
  else{
    return sz[r];
  }
}
int main(){
  int i;
  wmem = memarr;
  int ind;
  rd(N);
  {
    int FmcKpFmN;
    for(FmcKpFmN=(0);FmcKpFmN<(N);FmcKpFmN++){
      rd(C[FmcKpFmN]);C[FmcKpFmN] += (-1);
    }
  }
  {
    int t_ynMSdg;
    for(t_ynMSdg=(0);t_ynMSdg<(N-1);t_ynMSdg++){
      rd(A[t_ynMSdg]);A[t_ynMSdg] += (-1);
      rd(B[t_ynMSdg]);B[t_ynMSdg] += (-1);
    }
  }
  g.setEdge(N,N-1,A,B);
  hld.init(g);
  g.preorder(pre);
  for(i=(0);i<(N);i++){
    rev[pre[i]] = i;
  }
  g.SubTreeSize(0, sz);
  for(i=(0);i<(N);i++){
    b[C[pre[i]]].push_back(pre[i]);
  }
  for(i=(0);i<(N);i++){
    res[i] += (long long)N * (N+1) / 2;
    ind = 0;
    solve(-1, ind, b[i], res[i]);
  }
  {
    int jZyWAPpY;
    for(jZyWAPpY=(0);jZyWAPpY<(N);jZyWAPpY++){
      wt_L(res[jZyWAPpY]);
      wt_L('\n');
    }
  }
  return 0;
}
// cLay varsion 20200430-1 [beta]

// --- original code ---
// int N, A[2d5], B[2d5], C[2d5];
// graph g;
// HLD hld;
// int pre[2d5], prec[2d5], rev[2d5];
// 
// int ps;
// int sz[2d5];
// vector<int> b[2d5];
// 
// ll res[2d5];
// 
// int solve(int r, int &ind, vector<int> &b, ll &res){
//   int i, k = 0, nx;
//   ll s = 0, em;
// 
//   if(r==-1){
//     while(ind < b.size()){
//       i = ind++;
//       s += solve(b[i], ind, b, res);
//     }
//     em = N - s;
//     res -= em * (em+1) / 2;
//   } else {
//     rep(k,g.es[r]) if(g.edge[r][k] != hld.up(r)){
//       nx = g.edge[r][k];
//       while(ind < b.size() && hld.lca(nx, b[ind]) == nx){
//         i = ind++;
//         s += solve(b[i], ind, b, res);
//       }
//       em = sz[nx] - s;
//       res -= em * (em+1) / 2;
//       s = 0;
//     }
//   }
//   return if[r==-1, -1, sz[r]];
// }
// 
// {
//   int ind;
//   rd(N,(C--)(N),(A--,B--)(N-1));
//   g.setEdge(N,N-1,A,B);
//   hld.init(g);
// 
//   g.preorder(pre);
//   rep(i,N) rev[pre[i]] = i;
//   g.SubTreeSize(0, sz);
// 
//   rep(i,N) b[C[pre[i]]].push_back(pre[i]);
// 
//   rep(i,N){
//     res[i] += (ll)N * (N+1) / 2;
//     ind = 0;
//     solve(-1, ind, b[i], res[i]);
//   }
// 
//   wtLn(res(N));
// }


Compiler output:
main.c:2:9: fatal error: bits/stdc++.h: No such file or directory
    2 | #include<bits/stdc++.h>
      |         ^~~~~~~~~~~~~~~
compilation terminated.