P3258 [JLOI2014]松鼠的新家【树剖+线段树、树上差分两种解法】

P3258 [JLOI2014]松鼠的新家：https://www.luogu.com.cn/prob...
树剖+线段树

#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <string.h>
#define lson now<<1
#define rson now<<1|1

using namespace std;
const int N = 300010;
int a[N];
int n;
int head[N],nex[2*N],to[2*N],idx;
int u,v;
void add(int u,int v)
{
    to[idx]=v;
    nex[idx]=head[u];
    head[u]=idx++;
}
struct node
{
    long long int val;
    int lazy;
}tree[4*N];
int fa[N],deep[N],id[N],top[N],size[N],son[N];
int cnt;//dfs序编号
void dfs(int x,int f)//时间复杂度为O(2*m)
{
    fa[x]=f;
    size[x]=1;
    deep[x]=deep[f]+1;
    for(int i=head[x];~i;i=nex[i])
    {
        int j=to[i];
        if(j==f)
        {
            continue;
        }
        dfs(j,x);
        size[x]+=size[j];
        if(size[j]>size[son[x]])
        {
            son[x]=j;
        }
    }
}
void dfs2(int x,int t)//时间复杂度大约为O(2*m)
{
    top[x]=t;
    id[x]=++cnt;
    if(!son[x])
    {
        return;
    }
    dfs2(son[x],t);
    for(int i=head[x];~i;i=nex[i])
    {
        int j=to[i];
        if(j==fa[x]||j==son[x])
        {
            continue;
        }
        dfs2(j,j);
    }
}
void pushdown(int now,int l,int r)
{
    int ly=tree[now].lazy;
    tree[lson].lazy += ly;
    tree[rson].lazy += ly;
    int mid = (l+r)>>1;
    tree[lson].val += 1ll*(l-mid+1)*(1ll*ly);
    tree[rson].val += 1ll*(r-mid)*(1ll*ly);
    tree[now].lazy = 0;
}
void pushup(int now)
{
    tree[now].val = tree[lson].val + tree[rson].val;
}
void update(int now,int l,int r,int al,int ar,int val)
{
    if(al<=l&&ar>=r)
    {
        tree[now].val += 1ll*(r-l+1)*(1ll*val);
        tree[now].lazy += val;
        return;
    }
    if(tree[now].lazy!=0)
    {
        pushdown(now,l,r);
    }
    int mid = (l+r)>>1;
    if(al>mid)
    {
        update(rson,mid+1,r,al,ar,val);
    }
    else if(ar<=mid)
    {
        update(lson,l,mid,al,ar,val);
    }
    else
    {
        update(rson,mid+1,r,al,ar,val);
        update(lson,l,mid,al,ar,val);
    }
    pushup(now);
}

int query(int now,int l,int r,int al,int ar)
{
    if(al<=l&&ar>=r)
    {
        return tree[now].val;
    }
    if(tree[now].lazy!=0)
    {
        pushdown(now,l,r);
    }
    int mid = (l+r)>>1;
    long long int res = 0;
    if(al>mid)
    {
        res += query(rson,mid+1,r,al,ar);
    }
    else if(ar<=mid)
    {
        res += query(lson,l,mid,al,ar);
    }
    else
    {
        res += query(rson,mid+1,r,al,ar);
        res += query(lson,l,mid,al,ar);
    }
    return res;
}
void LCA(int x,int y)
{
    int fx=top[x],fy=top[y];
    while(fx!=fy)
    {
        if(deep[fx]>=deep[fy])
        {
            update(1,1,n,id[fx],id[x],1);
            x=fa[fx];
            fx=top[x];
        }
        else
        {
            update(1,1,n,id[fy],id[y],1);
            y=fa[fy];
            fy=top[y];
        }
    }
    if(deep[x]<deep[y])
    {
        swap(x,y);
    }
    update(1,1,n,id[y],id[x],1);
}

int main()
{
    scanf("%d",&n);
    memset(head,-1,sizeof(head));
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&a[i]);
    }
    for(int i=0;i<n-1;i++)
    {
        scanf("%d %d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1,0);
    dfs2(1,1);
    for(int i=1;i<n;i++)
    {
        LCA(a[i],a[i+1]);
    }
    for(int i=1;i<=n;i++)
    {
        long long ans = query(1,1,n,id[i],id[i]);
        if(i!=a[1])
        {
            ans --;
        }
        printf("%lld\n",ans);
    }
    return 0;
}

树上差分

#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <string.h>
#define lson now<<1
#define rson now<<1|1

using namespace std;
const int N = 300010;
int a[N];
int n;
int head[N],nex[2*N],to[2*N],idx;
int u,v;
void add(int u,int v)
{
    to[idx]=v;
    nex[idx]=head[u];
    head[u]=idx++;
}
int tmp[N];//存每个点在所有路径中出现的次数,不是上一种思想的存每个点到其父亲节点的边，看具体实际情况在这两种情况中切换

int fa[N][19],deep[N];
int cnt;//dfs序编号
void dfs(int x,int f)//时间复杂度为O(2*m)
{
    fa[x][0]=f;
    deep[x]=deep[f]+1;
    for(int i=1;i<=18;i++)
    {
        fa[x][i]=fa[fa[x][i-1]][i-1];
    }
    for(int i=head[x];~i;i=nex[i])
    {
        int j=to[i];
        if(j==f)
        {
            continue;
        }
        dfs(j,x);
    }
}
int get_lca(int x,int y)
{
    if(deep[x]<deep[y])
    {
        swap(x,y);
    }

    //第一步：先让x和y到同一高度
    for(int i=18;i>=0;i--)
    {
        if(deep[fa[x][i]]>=deep[y])
        {
            x=fa[x][i];
        }
    }
    if(x==y)//特殊情况，x跳到和y同一高度后重合，直接返回即可
    {
        return x;
    }
    //第二步：x和y不断往上跳，直到跳到最近公共祖先的下一层
    for(int i=18;i>=0;i--)
    {
        if(fa[x][i]!=fa[y][i])
        {
            x=fa[x][i];
            y=fa[y][i];
        }
    }
    return fa[x][0];
}
void dfs2(int x,int f)
{
    for(int i=head[x];~i;i=nex[i])
    {
        int j = to[i];
        if(j == f)
        {
            continue;
        }
        dfs2(j,x);
        tmp[x] += tmp[j];
    }
}

int main()
{
    scanf("%d",&n);
    memset(head,-1,sizeof(head));
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&a[i]);
    }
    for(int i=0;i<n-1;i++)
    {
        scanf("%d %d",&u,&v);
        add(u,v);
        add(v,u);
    }
    dfs(1,0);
    for(int i=1;i<n;i++)
    {
        tmp[a[i]]++;
        tmp[a[i+1]]++;
        int lca=get_lca(a[i],a[i+1]);
        tmp[lca]--;
        tmp[fa[lca][0]]--;
    }
    dfs2(1,0);
    for(int i=2;i<=n;i++)
    {
        tmp[a[i]]--;
    }
    for(int i=1;i<=n;i++)
    {
        printf("%d\n",tmp[i]);
    }
    return 0;
}