Bzoj3631: [JLOI2014]松鼠的新家

树剖后在dfn上差分

# include <bits/stdc++.h>
# define RG register
# define IL inline
# define Fill(a, b) memset(a, b, sizeof(a))
using namespace std;
typedef long long ll;
const int _(3e5 + 10);

IL ll Read(){
    RG ll x = 0, z = 1; RG char c = getchar();
    for(; c < '0' || c > '9'; c = getchar()) z = c == '-' ? -1 : 1;
    for(; c >= '0' && c <= '9'; c = getchar()) x = (x << 1) + (x << 3) + (c ^ 48);
    return x * z;
}

int n, fst[_], nxt[_ << 1], to[_ << 1], cnt, a[_], ans[_];
int dfn[_], Index, fa[_], size[_], top[_], son[_], deep[_];

IL void Add(RG int u, RG int v){  to[cnt] = v; nxt[cnt] = fst[u]; fst[u] = cnt++;  }

IL void Dfs1(RG int u, RG int ff){
    size[u] = 1; fa[u] = ff; deep[u] = deep[ff] + 1;
    for(RG int e = fst[u]; e != -1; e = nxt[e]){
        if(to[e] == ff) continue;
        Dfs1(to[e], u);
        size[u] += size[to[e]];
        if(size[to[e]] > size[son[u]]) son[u] = to[e];
    }
}

IL void Dfs2(RG int u, RG int Top){
    top[u] = Top; dfn[u] = ++Index;
    if(son[u]) Dfs2(son[u], Top);
    for(RG int e = fst[u]; e != -1; e = nxt[e]) if(!dfn[to[e]]) Dfs2(to[e], to[e]);
}

IL void Cover(RG int u, RG int v){
    while(top[u] ^ top[v]){
        if(deep[top[u]] < deep[top[v]]) swap(u, v);
		++ans[dfn[top[u]]]; --ans[dfn[u] + 1];
        u = fa[top[u]];
    }
    if(deep[u] > deep[v]) swap(u, v);
	++ans[dfn[u]]; --ans[dfn[v] + 1];
}

int main(RG int argc, RG char* argv[]){
    n = Read();
	for(RG int i = 1; i <= n; ++i) a[i] = Read(), fst[i] = -1;
	for(RG int i = 1, x, y; i < n; ++i) x = Read(), y = Read(), Add(x, y), Add(y, x);
	Dfs1(1, 0); Dfs2(1, 1);
	for(RG int i = 1; i < n; ++i) Cover(a[i], a[i + 1]);
	for(RG int i = 1; i <= n; ++i) ans[i] += ans[i - 1];
	for(RG int i = 1; i <= n; ++i) printf("%d\n", ans[dfn[i]] - (i != a[1]));
    return 0;
}