lca + 差分模板：R90284800 记录详情

题目链接：

P3258 [JLOI2014]松鼠的新家 - 洛谷 | 计算机科学教育新生态 (luogu.com.cn)

#include <queue>
#include <cstdio>
#include <cstring>
#include <iostream>

using namespace std;

const int N = 300010, M = 2 * N;

int n;
int w[N];
int d[N];
int fa[N][20], depth[N];
int h[N], e[M], ne[M], idx;

void add(int a, int b)//邻接表模板
{
    e[idx] = b;
    ne[idx] = h[a];
    h[a] = idx;
    idx ++ ;
}

void bfs(int u, int pa)//预处理fa, depth模板
{
    queue<int> q;
    q.push(u);
    memset(depth, 0x3f, sizeof depth);
    depth[0] = 0, depth[1] = 1;
    
    while (q.size())
    {
        int t = q.front();
        q.pop();
        
        for (int i = h[t]; i != -1; i = ne[i])
        {
            int j = e[i];
            if (j == pa) continue;
            if (depth[j] > depth[t] + 1)
            {
                q.push(j);
                depth[j] = depth[t] + 1;
                fa[j][0] = t;
                for (int k = 1; k <= 19; k ++ )
                    fa[j][k] = fa[fa[j][k - 1]][k - 1];
            }
        }
    }
}

int lca(int a, int b)//最近公共祖先模板
{
    if (a == b) return 0;
    if (depth[a] < depth[b]) swap(a, b);
    
    for (int k = 19; k >= 0; k -- )
        if (depth[fa[a][k]] >= depth[b])
            a = fa[a][k];
            
    if (a == b) return a;
    
    for (int k = 19; k >= 0; k -- )
        if (fa[a][k] != fa[b][k])
        {
            a = fa[a][k];
            b = fa[b][k];
        }
        
    return fa[a][0];
}

void dfs(int u, int fa)//求出差分数组的和
{
    for (int i = h[u]; i != -1; i = ne[i])
    {
        int j = e[i];
        if (j == fa) continue;
        dfs(j, u);
        d[u] += d[j];
    }
}

int main()
{
    cin >> n;
    memset(h, -1, sizeof h);
    for (int i = 1; i <= n; i ++ ) scanf("%d", &w[i]);
    
    for (int i = 0; i < n - 1; i ++ )
    {
        int a, b;
        scanf("%d %d", &a, &b);
        add(a, b), add(b, a);
    }
    
    bfs(1, 0);//预处理fa, 与 depth数组
    
    for (int i = 1; i <= n - 1; i ++ )//计算每个点之间的差分数组
    {
        int a = w[i], b = w[i + 1];
        int anc = lca(a, b);
        d[a] ++ , d[b] ++ , d[anc] -- , d[fa[anc][0]] -- ;
    }
    
    dfs(1, 0);
    
    for (int i = 2; i <= n; i ++ )//除了开始节点和结束节点之外每个点都被加了2次。而结束节点根据题意不需要算
        d[w[i]] -- ;
    
    for (int i = 1; i <= n; i ++ ) printf("%d\n", d[i]);
    
    return 0;
}