树上差分

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 300000;
int n;
int a[MAXN + 5];
vector<int> e[MAXN + 5];
// -------start LCA--------
// f[i][j] 记录 i 的 2^j 级别祖先
int f[MAXN + 5][25];
// dep[i] 求节点 i 的深度
int dep[MAXN + 5];
void dfs(int u, int fa)
{
    f[u][0] = fa;
    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i];
        if (v == fa)
            continue;
        dep[v] = dep[u] + 1;
        dfs(v, u);
    }
}
void init(int s)
{
    dep[s] = 0;
    dfs(s, 0); // 预处理深度 dep[u]，以及一级祖先 f[u][0]
    // i 的 2^j 级别祖先 = i 的 2^{j-1} 级别祖先的 2^{j-1} 级别祖先
    for (int j = 1; (1LL << j) <= n; j++)
        for (int i = 1; i <= n; i++)
            f[i][j] = f[f[i][j - 1]][j - 1];
}
// lca(u,v) 返回 u,v 的最近公共祖先
int lca(int u, int v)
{
    // 保证 u 在上面，v 在下面
    if (dep[v] < dep[u])
        swap(u, v);
    // 拉到同样的深度
    for (int j = 20; j >= 0; j--)
        if (dep[v] - dep[u] >= (1 << j))
            v = f[v][j];
    // 初始两点之间是祖孙关系时，拉到同样深度就会变成同一个点
    if (u == v)
        return u;
    // 同步往上走，跳到了 lca 下面一跳位
    for (int j = 20; j >= 0; j--)
        if (f[u][j] != f[v][j])
            u = f[u][j], v = f[v][j];
    return f[u][0];
}
// -------end LCA--------
// 树上差分数组
// 每个节点的真实权值为子树所有节点的 d 之和
int d[MAXN + 5];
void dfsSum(int u, int fa)
{
    // sum[u] = d[u];
    for (int i = 0; i < e[u].size(); i++)
    {
        int v = e[u][i];
        if (v == fa)
            continue;
        dfsSum(v, u);
        d[u] += d[v];
        // sum[u]+=sum[v];
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n;
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1; i <= n - 1; i++)
    {
        int u, v;
        cin >> u >> v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    // ----- start LCA 预处理 -----
    init(1);
    // ----- end LCA 预处理 -----
    for (int i = 1; i <= n - 1; i++)
    {
        int u, v;
        u = a[i];
        v = a[i + 1];
        // [u, v) 路径 +1
        // 先对 [u, v] 路径 +1
        int uv = lca(u, v);
        d[u]++, d[v]++, d[uv]--, d[f[uv][0]]--;
        // 对 v 单点 -1
        d[v]--, d[f[v][0]]++;
    }
    dfsSum(1, 0);
    for (int i = 1; i <= n; i++)
        cout << d[i] << "\n";
    return 0;
}