P3258 [JLOI2014] 树上差分

题意

传送门 P3258 [JLOI2014]松鼠的新家

题解

考虑任意两个房间之间的路径，若为最后一条路径，则终点不准备糖果；反之，其终点必为下一条路径的起点。那么可以将路径处理为左闭右开 $[a_i,a_{i+1})$。树上差分，然后求前缀和即可。开边界不易处理，那么将路径处理为 $[a_i,a_{i+1}]$，同时记录终点所在节点需要在前缀和中减一。

倍增法求解 $LCA$，总时间复杂度 $O(N\log N)$

#include <bits/stdc++.h>
using namespace std;
const int maxn = 300005, maxlg = 19;
int N, A[maxn], cnt[maxn], del[maxn];
int lg[maxn], par[maxn][maxlg], dep[maxn];
vector<int> G[maxn];

void dfs(int u, int p, int d)
{
    
    
    par[u][0] = p, dep[u] = d;
    for (int i = 1; i <= lg[dep[u]]; ++i)
        par[u][i] = par[par[u][i - 1]][i - 1];
    for (auto &v : G[u])
        if (v != p)
            dfs(v, u, d + 1);
}

int lca(int u, int v)
{
    
    
    if (dep[u] < dep[v])
        swap(u, v);
    while (dep[u] > dep[v])
        u = par[u][lg[dep[u] - dep[v]]];
    if (u == v)
        return u;
    for (int k = lg[dep[u]]; k >= 0;)
        if (par[u][k] != par[v][k])
            u = par[u][k], v = par[v][k], k = lg[dep[u]];
        else
            --k;
    return par[u][0];
}

void get_sum(int u, int p)
{
    
    
    for (auto &v : G[u])
        if (v != p)
            get_sum(v, u), cnt[u] += cnt[v];
}

int main()
{
    
    
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> N;
    for (int i = 0; i < N; ++i)
        cin >> A[i], --A[i];
    lg[0] = -1;
    for (int i = 1; i < N; ++i)
        lg[i] = lg[i >> 1] + 1;
    for (int i = 1; i < N; ++i)
    {
    
    
        int u, v;
        cin >> u >> v;
        --u, --v;
        G[u].push_back(v), G[v].push_back(u);
    }
    dfs(0, -1, 0);
    for (int i = 0; i + 1 < N; ++i)
    {
    
    
        int u = A[i], v = A[i + 1];
        ++del[v];
        int p = lca(u, v);
        ++cnt[u], ++cnt[v], --cnt[p];
        if (dep[p])
            --cnt[par[p][0]];
    }
    get_sum(0, -1);
    for (int i = 0; i < N; ++i)
        cout << cnt[i] - del[i] << '\n';
    return 0;
}