luogu3258 [JLOI2014]松鼠的新家

#include <iostream>
#include <cstdio>
using namespace std;
int n, uu, vv, a[300005], hea[300005], dep[300005], gnd[300005][21], qwq;
int siz[300005], son[300005], top[300005], cnt, idx[300005], ans[300005];
struct Edge{
    int too, nxt;
}edge[600005];
struct SGT{
    int sum[1200005];
    int tag[1200005];
    void pushDown(int o, int l, int r, int lson, int rson, int mid){
        tag[lson] += tag[o];
        tag[rson] += tag[o];
        sum[lson] += (mid-l+1) * tag[o];
        sum[rson] += (r-mid) * tag[o];
        tag[o] = 0;
    }
    void update(int o, int l, int r, int x, int y, int k){
        if(l>=x && r<=y){
            sum[o] += (r-l+1) * k;
            tag[o] += k;
        }
        else{
            int mid=(l+r)>>1;
            int lson=o<<1;
            int rson=lson|1;
            if(tag[o])  pushDown(o, l, r, lson, rson, mid);
            if(x<=mid)  update(lson, l, mid, x, y, k);
            if(mid<y)   update(rson, mid+1, r, x, y, k);
            sum[o] = sum[lson] + sum[rson];
        }
    }
    void func(int o, int l, int r){
        if(l==r)    ans[l] += sum[o];
        else{
            int mid=(l+r)>>1;
            int lson=o<<1;
            int rson=lson|1;
            if(tag[o])  pushDown(o, l, r, lson, rson, mid);
            if(l<=mid)  func(lson, l, mid);
            if(mid<r)   func(rson, mid+1, r);
        }
    }
}sgt;
void add_edge(int fro, int too){
    edge[++cnt].nxt = hea[fro];
    edge[cnt].too = too;
    hea[fro] = cnt;
}
void dfs1(int x, int f){
    gnd[x][0] = f;
    dep[x] = dep[f] + 1;
    siz[x] = 1;
    int maxSon=-1;
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(t==f)    continue;
        dfs1(t, x);
        siz[x] += siz[t];
        if(siz[t]>maxSon){
            maxSon = siz[t];
            son[x] = t;
        }
    }
}
void dfs2(int x, int topf){
    top[x] = topf;
    idx[x] = ++qwq;
    if(!son[x]) return ;
    dfs2(son[x], topf);
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(t==gnd[x][0] || t==son[x])   continue;
        dfs2(t, t);
    }
}
int getLca(int xx, int yy){
    if(dep[xx]<dep[yy]) swap(xx, yy);
    for(int i=19; i>=0; i--)
        if(dep[gnd[xx][i]]>=dep[yy])
            xx = gnd[xx][i];
    if(xx==yy)  return xx;
    for(int i=19; i>=0; i--)
        if(gnd[xx][i]!=gnd[yy][i]){
            xx = gnd[xx][i];
            yy = gnd[yy][i];
        }
    return gnd[xx][0];
}
void updRange(int uu, int vv){
    while(top[uu]!=top[vv]){
        if(dep[top[uu]]<dep[top[vv]])   swap(uu, vv);
        sgt.update(1, 1, n, idx[top[uu]], idx[uu], 1);
        uu = gnd[top[uu]][0];
    }
    if(dep[uu]>dep[vv]) swap(uu, vv);
    sgt.update(1, 1, n, idx[uu], idx[vv], 1);
}
int main(){
    cin>>n;
    for(int i=1; i<=n; i++)
        scanf("%d", &a[i]);
    for(int i=1; i<n; i++){
        scanf("%d %d", &uu, &vv);
        add_edge(uu, vv);
        add_edge(vv, uu);
    }
    dep[1] = 1;
    dfs1(1, 0);
    dfs2(1, 1);
    for(int i=1; i<=19; i++)
        for(int j=1; j<=n; j++)
            gnd[j][i] = gnd[gnd[j][i-1]][i-1];
    for(int i=1; i<n; i++){
        uu = a[i];
        vv = a[i+1];
        int lca=getLca(uu, vv);
        ans[idx[lca]]--;
        updRange(uu, lca);
        updRange(vv, lca);
    }
    sgt.func(1, 1, n);
    for(int i=2; i<=n; i++)
        ans[idx[a[i]]]--;
    for(int i=1; i<=n; i++)
        printf("%d\n", ans[idx[i]]);
    return 0;
}