牛客 - 阔力梯的树(树上启发式合并)

题目链接:点击查看

题目大意:给出一棵树,每个节点都有一个编号,现在规定每个节点的“结实程度”为其子树中所有的编号排序后相邻两个数之差的平方和,现在需要求出每个节点的“结实程度”

题目分析:因为离线+对每个子树的操作,所以可以用树上启发式合并nlogn解决问题, 对于每次更新/删除操作,因为题目的计算方法的前提是有序,并且每个编号两两都不互相重复,所以我们可以直接用一个set来维护,对于统计答案,我一开始用了一个很笨的方法,对于每个节点都遍历一遍set内的数,这样时间复杂度显然是n*n*logn的,肯定超时了,其实借助树上启发式合并不断传递的思想,对于每次更新的一个数,我们都可以在logn的时间复杂度内解决,具体实现就是先二分找到这个操作数的位置,并取出其相邻的两个数,减去之前的贡献,累加上新增加的贡献即可完成贡献,这样总时间复杂度为nlognlogn,而且牛客的评测机跑的好快,我用vector邻接表存边最后也才跑了300多ms

在代码实现方面,关于对一个数维护答案的代码部分,增加时和删除时的代码是大同小异的,但我为了方便直接CV大法了,没有考虑到代码复用的问题,可能代码看起来比较冗杂

代码:

#include<iostream>
#include<cstdio> 
#include<string>
#include<ctime>
#include<cstring>
#include<algorithm>
#include<stack>
#include<queue>
#include<map>
#include<set>
#include<cmath>
#include<sstream>
#include<unordered_map>
using namespace std;
 
typedef long long LL;
 
const int inf=0x3f3f3f3f;
 
const int N=1e5+100;

vector<int>node[N];

set<int>st;

LL sum;

int son[N],num[N];

LL ans[N];

bool vis[N];

void dfs_son(int u)
{
	son[u]=-1;
	num[u]=1;
	for(auto v:node[u])
	{
		dfs_son(v);
		num[u]+=num[v];
		if(son[u]==-1||num[son[u]]<num[v])
			son[u]=v;
	}
}

void add(int num)
{
	set<int>::iterator it=st.lower_bound(num),start=st.begin(),end=st.end();
	end--;
	auto pre=it,next=it;
	if(it!=start)
		pre--;
	if(it!=end)
		next++;
	if(pre!=it&&next!=it)//当前插入的数在中间
	{
		sum-=1LL*(*next-*pre)*(*next-*pre);
		sum+=1LL*(*next-*it)*(*next-*it);
		sum+=1LL*(*pre-*it)*(*pre-*it);
	} 
	else if(pre!=it)//在最后 
	{
		sum+=1LL*(*pre-*it)*(*pre-*it);
	}
	else if(next!=it)//在开头
	{
		sum+=1LL*(*next-*it)*(*next-*it);
	}
	else//只有一个数 
	{
		sum=0;
	}
}

void del(int num)
{
	set<int>::iterator it=st.lower_bound(num),start=st.begin(),end=st.end();
	end--;
	auto pre=it,next=it;
	if(it!=start)
		pre--;
	if(it!=end)
		next++;
	if(pre!=it&&next!=it)//当前删除的数在中间
	{
		sum+=1LL*(*next-*pre)*(*next-*pre);
		sum-=1LL*(*next-*it)*(*next-*it);
		sum-=1LL*(*pre-*it)*(*pre-*it);
	} 
	else if(pre!=it)//在最后 
	{
		sum-=1LL*(*pre-*it)*(*pre-*it);
	}
	else if(next!=it)//在开头
	{
		sum-=1LL*(*next-*it)*(*next-*it);
	}
	else//只有一个数 
	{
		sum=0;
	}
}

void update(int u,int val)
{
	if(val==1)
	{
		st.insert(u);
		add(u);
	}
	else
	{
		del(u);
		st.erase(u);
	}
	for(auto v:node[u])
	{
		if(vis[v])
			continue;
		update(v,val);
	}
}

void dfs(int u,int keep)
{
	for(auto v:node[u])
	{
		if(v==son[u])
			continue;
		dfs(v,0);
	}
	if(son[u]!=-1)
	{
		dfs(son[u],1);
		vis[son[u]]=true;
	}
	update(u,1);
	ans[u]=sum;
	if(son[u]!=-1)
		vis[son[u]]=false;
	if(!keep)
		update(u,-1);
}

int main()
{
//	freopen("input.txt","r",stdin);
//	ios::sync_with_stdio(false);
	int n;
	scanf("%d",&n);
	for(int i=2;i<=n;i++)
 	{
 		int fa;
 		scanf("%d",&fa);
 		node[fa].push_back(i);
	}
	dfs_son(1);
	dfs(1,1);
 	for(int i=1;i<=n;i++)
 		printf("%lld\n",ans[i]);
 
 
 
	
	
	
	
	
	
	
	
	
	return 0;
}
发布了549 篇原创文章 · 获赞 16 · 访问量 1万+

猜你喜欢

转载自blog.csdn.net/qq_45458915/article/details/104071312