U95602 射手座之日

· · 个人记录

树上启发式合并计算

最开始ans[u]因该是覆盖这颗子树的数量 最终答案应该等于ans[u]减去每颗子树的ans[t]就是覆盖这颗子树的根节点的区间数量;

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
#include<set>
using namespace std;

typedef long long LL;
typedef pair<int,int> PII;
const int N=2e5+10,mod=998244353;

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

int son[N],siz[N];
LL ans[N];

int l[N],r[N],id[N];
int find(int fa[],int u){if(u!=fa[u])fa[u]=find(fa,fa[u]);return fa[u];}
LL find(int t){l[t]=t-1;r[t]=t+1;return LL(t-find(l,t))*(find(r,t)-t);}

void dfs(int u)
{
    siz[u]=1; 
    for(int i=h[u];~i;i=ne[i])
    {
        int t=e[i];
        dfs(t);
        if(siz[son[u]]<siz[t])son[u]=t;
        siz[u]+=siz[t];
    }
    return ; 
}

void calc(int u,int val,LL &x)
{   
    int t=id[u];

    if(val==-1) l[t]=r[t]=t;
    else x+=find(t);

    for(int i=h[u];~i;i=ne[i])
        calc(e[i],val,x);
}

void dfs(int u,int keep)
{
    for(int i=h[u];~i;i=ne[i])
    {
        int t=e[i];
        if(t==son[u])continue;
        dfs(t,0);
    }
    if(son[u])dfs(son[u],1);

    ans[u]+=find(id[u]);
    for(int i=h[u];~i;i=ne[i])
    {
        int t=e[i];
        if(t==son[u]){
            ans[u]+=ans[t];
            continue;
        }
        calc(t,1,ans[u]);
    }

    if(!keep)calc(u,-1,ans[u]);
}

int main()
{
    int n;
    scanf("%d",&n);
    for(int i=0;i<=n+1;i++)l[i]=r[i]=i,h[i]=-1;

    for(int i=2;i<=n;i++){int a;scanf("%d",&a);add(h,a,i);}
    for(int i=1;i<=n;i++){int a;scanf("%d",&a);id[a]=i;}

    dfs(1); dfs(1,1);
    LL X=0;

    for(int I=1;I<=n;I++)
    {
        int w;
        LL T=ans[I];
        scanf("%d",&w);
        for(int i=h[I];~i;i=ne[i])
        {
            int t=e[i];
            T-=ans[t];
        }
        X+=T*w;
    }
    cout<<X<<endl;
    return 0;
}