原题链接 松鼠的新家

题目描述

给出一棵树和一个长度为 $n$ 的序列 $p$，小熊要从 $p_1$ 走到 $p_2$，再从 $p_2$ 走到 $p_3$，以此类推走到 $p_n$
小熊每抵达一个 $p_n$ 之外的地方都要吃一块糖，求每个点要准备多少块糖
$n \leq 3 × 10^5$

输入样例

5
1 4 5 3 2
1 2
2 4
2 3
4 5

输出样例

1
2
1
2
1

树上差分

首先我们看出小熊从 $a$ 点走到 $b$ 点，差不多都要给 $a - b$ 路径上的点 + 1

而从 $a$ 到 $b$ 再从 $b$ 到 $c$， $b$ 是不用被算两次的，结合题目说最终的终点不需要计算

所以可以把小熊每一次从 $a$ 点走到 $b$ 点行动，看做将 $a - b$ 路径上除了点b之外的点 + 1

判断一下路径的形状分类讨论即可

时间复杂度 $O(nlogn)$

C++ 代码

#include<cstdio>
#include<cstring>
#include<algorithm>

using namespace std;

const int N = 300010, M = N << 1;

int h[N], e[M], ne[M], idx;
void add(int a, int b)
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
}

int n;
int p[N];

int f[N][20], dep[N];
void dfs(int u, int fa)
{
    dep[u] = dep[fa] + 1;
    f[u][0] = fa;
    for(int i = 1; i < 20; i ++)
        f[u][i] = f[f[u][i - 1]][i - 1];

    for(int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if(j == fa)  continue;
        dfs(j, u);
    }
}

int lca(int a, int b)
{
    if(dep[a] < dep[b])  swap(a, b);

    for(int i = 19; i >= 0; i --)
        if(dep[f[a][i]] >= dep[b])
            a = f[a][i];

    if(a == b)  return a;

    for(int i = 19; i >= 0; i --)
        if(f[a][i] != f[b][i])
            a = f[a][i], b = f[b][i];
    return f[a][0];
}

int res[N];

void dfs2(int u)
{
    for(int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if(j == f[u][0])  continue;
        dfs2(j);
        res[u] += res[j];
    }
}

int main()
{
    memset(h, -1, sizeof h);

    scanf("%d", &n);
    for(int i = 1; i <= n; i ++)  scanf("%d", &p[i]);

    for(int i = 1; i < n; i ++)
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }

    dfs(1, 0);

    for(int i = 1; i < n; i ++)
    {
        int a = p[i], b = p[i + 1];

        int t = lca(a, b);

        //printf("a = %d b = %d lca = %d\n", a, b, t);
        if(t == b)  res[a] ++, res[b] --;
        else if(t == a)  res[f[a][0]] --, res[f[b][0]] ++;
        else  res[a] ++, res[f[b][0]] ++, res[t] --, res[f[t][0]] --;
    }

    dfs2(1);

    for(int i = 1; i <= n; i ++)  printf("%d\n", res[i]);
    return 0;
}