2019 ICPC Asia Nanchang Regional K.Tree 树上启发式合并 + 动态开点线段树

传送门

文章目录

题意:

给你一棵树,每个点都有一个权值 v a l val val,求满足以下条件
(1) x ! = y x!=y x!=y
(2) x x x y y y不互为祖先
(3) v a l [ l c a ( x , y ) ] ∗ 2 = v a l [ x ] + v a l [ y ] val[lca(x,y)]*2=val[x]+val[y] val[lca(x,y)]2=val[x]+val[y]
(4) l e n ( x , y ) < = k len(x,y)<=k len(x,y)<=k
( x , y ) (x,y) (x,y)的对数。

思路:

可以发现这个点对是可以转换成一颗子树的问题来递归解决的,他们以这个子树的根为 l c a lca lca,让后两个点一定存在于子树的不同分支里面,碰到子树问题我们自然的想到用 d s u dsu dsu来解决。
当我们要更新一颗子树的时候,当前子树的重儿子的信息已经存起来了,我们只需要依次遍历计算其他分支,让后更新其他分支的信息。要注意必须先遍历计算才能更新,这样才能保证算出来的一定是根的不同分支里的对数。让后还有一个问题,就是怎么算答案呢?我们更新的时候是将 d e p t h [ u ] depth[u] depth[u] v a l [ u ] val[u] val[u]一起放到要更新的数组里面的,我们查询的时候要查询权值为 v a l [ r t ] ∗ 2 − v a l [ n o w ] val[rt]*2-val[now] val[rt]2val[now]且深度 < = d e p t h [ r t ] + k − ( d e p t h [ n o w ] − d e p t h [ r t ] ) <=depth[rt]+k-(depth[now]-depth[rt]) <=depth[rt]+k(depth[now]depth[rt])的点对。显然普通数据结构不好维护,所以我们可以对每个权值开一颗权值线段树,让后动态开点就好啦,就是有点不好调。
最后答案乘二即可。

//#pragma GCC optimize(2)
#include<cstdio>
#include<iostream>
#include<string>
#include<cstring>
#include<map>
#include<cmath>
#include<cctype>
#include<vector>
#include<set>
#include<queue>
#include<algorithm>
#include<sstream>
#include<ctime>
#include<cstdlib>
#define X first
#define Y second
#define L (u<<1)
#define R (u<<1|1)
#define pb push_back
#define mk make_pair
#define Mid (tr[u].l+tr[u].r>>1)
#define Len(u) (tr[u].r-tr[u].l+1)
#define random(a,b) ((a)+rand()%((b)-(a)+1))
#define db puts("---")
using namespace std;

//void rd_cre() { freopen("d://dp//data.txt","w",stdout); srand(time(NULL)); }
//void rd_ac() { freopen("d://dp//data.txt","r",stdin); freopen("d://dp//AC.txt","w",stdout); }
//void rd_wa() { freopen("d://dp//data.txt","r",stdin); freopen("d://dp//WA.txt","w",stdout); }

typedef long long LL;
typedef unsigned long long ULL;
typedef pair<LL,LL> PII;

const int N=200010,mod=1e9+7,INF=0x3f3f3f3f;
const double eps=1e-6;

int n,k;
vector<int>v[N];
int val[N];
LL cnt,ans[N];
int depth[N],se[N],son[N];
int root[N],idx;
struct Node
{
    
    
    int l,r;
    int cnt;
}tr[N*40];

void modify(int &rt,int l,int r,int x,int tag)
{
    
    
    if(!rt) rt=++idx;
    if(l==r) {
    
     tr[rt].cnt+=tag; return; }
    int mid=l+r>>1;
    if(x<=mid) modify(tr[rt].l,l,mid,x,tag);
    else modify(tr[rt].r,mid+1,r,x,tag);
    tr[rt].cnt=tr[tr[rt].l].cnt+tr[tr[rt].r].cnt;
}

int query(int rt,int l,int r,int ql,int qr)
{
    
    
    if(!rt) return 0;
    if(l>=ql&&r<=qr) return tr[rt].cnt;
    int mid=l+r>>1;
    int ans=0;
    if(ql<=mid) ans+=query(tr[rt].l,l,mid,ql,qr);
    if(qr>mid) ans+=query(tr[rt].r,mid+1,r,ql,qr);
    return ans;
}

void dfs_son(int u,int fa)
{
    
    
    se[u]=1; depth[u]=depth[fa]+1;
    for(auto x:v[u])
    {
    
    
        if(x==fa) continue;
        dfs_son(x,u);
        se[u]+=se[x];
        if(se[x]>se[son[u]]) son[u]=x;
    }
}

void update(int u,int fa,int son,int tag)
{
    
    
    modify(root[val[u]],1,n,depth[u],tag);
    for(auto x:v[u])
    {
    
    
        if(x==fa||x==son) continue;
        update(x,u,son,tag);
    }
}

void solve(int u,int fa,int son,int rt)
{
    
    
    if(val[rt]*2-val[u]>=0&&val[rt]*2-val[u]<=n&&depth[rt]+k-(depth[u]-depth[rt])>=depth[rt]+1) cnt+=query(root[val[rt]*2-val[u]],1,n,depth[rt]+1,min(depth[rt]+k-(depth[u]-depth[rt]),n));
    //cnt+=mp[{depth[rt]+k-(depth[u]-depth[rt]),val[rt]*2-val[u]}];
    for(auto x:v[u])
    {
    
    
        if(x==fa||x==son) continue;
        solve(x,u,son,rt);
    }
}

void dfs(int u,int fa,int op)
{
    
    
    for(auto x:v[u])
    {
    
    
        if(x==fa||x==son[u]) continue;
        dfs(x,u,0);
    }
    if(son[u]) dfs(son[u],u,1);
    for(auto x:v[u])
    {
    
    
        if(x==fa||x==son[u]) continue;
        solve(x,u,son[u],u); update(x,u,son[u],1);
    }
    modify(root[val[u]],1,n,depth[u],1);
    ans[u]=cnt; cnt=0;
    if(!op) update(u,fa,-1,-1);
}

int main()
{
    
    
//	ios::sync_with_stdio(false);
//	cin.tie(0);

    scanf("%d%d",&n,&k);
    for(int i=1;i<=n;i++) scanf("%d",&val[i]);
    for(int i=2;i<=n;i++)
    {
    
    
        int x; scanf("%d",&x);
        v[x].pb(i);
    }
    dfs_son(1,0);
    dfs(1,0,1);
    LL sum=0;
    for(int i=1;i<=n;i++) sum+=ans[i];
    printf("%lld\n",sum*2);


	return 0;
}
/*
4 3
1 1 3 2
1 1 2
*/




猜你喜欢

转载自blog.csdn.net/m0_51068403/article/details/115192427
今日推荐