【JLOI 2014】 松鼠的新家

【题目链接】

           https://www.lydsy.com/JudgeOnline/problem.php?id=3631

【算法】

           树上差分

【代码】

             

#include<bits/stdc++.h>
using namespace std;
#define MAXN 300010
#define MAXLOG 20

int i,n,u,v,Lca;
int a[MAXN],dep[MAXN],cnt[MAXN];
int anc[MAXN][MAXLOG];
vector< int > e[MAXN];

inline int lca(int u,int v)
{
        int i,t;
        if (dep[u] > dep[v]) swap(u,v);
        t = dep[v] - dep[u];
        for (i = 0; i < MAXLOG; i++)
        {
                if (t & (1 << i))
                        v = anc[v][i];
        }
        if (u == v) return u;
        for (i = MAXLOG - 1; i >= 0; i--)
        {
                if (anc[u][i] != anc[v][i])
                {
                        u = anc[u][i];
                        v = anc[v][i];
                }
        }
        return anc[u][0];
}
inline void dfs(int u)
{
        int i;
        for (i = 1; i < MAXLOG; i++) anc[u][i] = anc[anc[u][i-1]][i-1];
        for (i = 0; i < e[u].size(); i++)
        {
                if (anc[u][0] != e[u][i])
                {
                        dep[e[u][i]] = dep[u] + 1;
                        anc[e[u][i]][0] = u;
                        dfs(e[u][i]);        
                }    
        }        
}
inline void calc(int u)
{
        int i;
        for (i = 0; i < e[u].size(); i++)
        {
                if (anc[u][0] != e[u][i])
                {
                        calc(e[u][i]);
                        cnt[u] += cnt[e[u][i]];
                }
        }
}

int main() {
        
        scanf("%d",&n);
        for (i = 1; i <= n; i++) scanf("%d",&a[i]);
        for (i = 1; i < n; i++)
        {
                scanf("%d%d",&u,&v);
                e[u].push_back(v);
                e[v].push_back(u);
        }
        dfs(1);
        for (i = 1; i < n; i++)
        {
                Lca = lca(a[i],a[i+1]);        
                cnt[a[i]]++;
                cnt[a[i+1]]++;
                cnt[Lca]--;
                cnt[anc[Lca][0]]--;
        }
        calc(1);
        for (i = 2; i <= n; i++) cnt[a[i]]--;
        for (i = 1; i <= n; i++) printf("%d\n",cnt[i]);
        
        return 0;
    
}