ABC359

坐标亚洲最大的中学，三轮省集晚自习时候打的。

下午打了线段树合并和一些莫队。

G

哥们，懂不懂赛时过 G 的含金量啊！

歌词一眼秒了，然后一直在贬低 AT，于是来看 G 了。

首先按照颜色分类，设目前考虑的是第 $t$ 类，点集为 $S_{t}$ 。

要求 $S_{t}$ 的完全图的边权和。将 $d i s (x, y)$ 转换为 $d i s (x, l c a) + d i s (l c a, y)$

建出虚树，设虚树的点集为 $S_{t}^{'}$ ， $S_{t} \in S_{t}^{'}$ 。考虑所有 $x \in S_{t}^{'}$ ，将 $x$ 视作 lca，那么，两个点 $a, b \in S_{t}$ 以 $x$ 为其 lca 当且仅当 $a, b$ 在 $x$ 的不同子树中。直接枚举是 $O (n^{3} \log n)$ 的，还不如暴力，暴力利用倍增可以做到 $O (n^{2} \log n)$ 。

变一下，假设我们知道了 $x, a$ ，如何求出所有的 $\sum d i s (a, b)$ ？设 $a$ 所在子树的根为 $r t_{a}$ ，该子树中的与 $a$ 同类的点有 $n u m_{a}$ 个， $x$ 为根的子树中同类点有 $n u m_{x}$ 个。则 $b$ 的数量为 $n u m_{x} - n u m_{a}$ 。

这样 $d i s (l c a, b)$ 还是不好求，我们考虑只求 $d i s (a, l c a)$ 。

求一个 $d i s (a, l c a)$ 是不难的，我们考虑要求几个。显然是 $n u m_{x} - n u m_{a}$ 个。于是这样可以做到 $O (n^{2} \log n)$ 。

能不能再去掉一个 $n$ ？

设 $r t_{a}$ 为根的子树中同类点的集合为 $S_{a}$ 。我们要求 $\sum_{z \in S_{a}} d i s (z, l c a) \times (n u m_{x} - n u m_{a})$ ，直接提出来，维护 $s u m$ 代表 $\sum_{z \in S_{a}} d i s (z, l c a)$ ，这个在 dfs 时可以处理出来。

需要的几个工具：快速求 lca，快速求树上两点间距离，虚树。

三个各需要一个 dfs。

点击查看代码

 #include<bits/stdc++.h>
#define int long long
#define fi first
#define se second
 
//#define file(...) freopen("##__VA_ARGS__.in","r",stdin);freopen("##__VA_ARGS__.out","w",stdout)
//#define outfile(...) freopen("##__VA_ARGS__.out","w",stdout)
//#define infile(...) freopen("##__VA_ARGS__.in","r",stdin);
//#define datafile(...) freopen("##__VA_ARGS__.in","w",stdout)
#define debug(...) fprintf(stderr,##__VA_ARGS__)
//#define printf(...) fprintf(stdout,##__VA_ARGS)
template<typename T,typename I>
#define heap(T) std::priority_queue<T>
#define umap(T,I) std::unordered_map<T,I>
#define pii(T,I) std::pair<T,I>
#define map(T,I) std::map<T,I>
#define vector(T) std::vector<T>
#define stack(T) std::stack<T>
using pii=std::pair<int,int>;
using ll=long long;
using i128=__int128;
using ull=unsigned long long;
 
template<typename T>
inline void read(T &x){
	x=0;
	char c=getchar();
	int f=1;
	while(c<'0'||c>'9'){
		if(c=='-') f=-1;
		c=getchar();
	}
	while(c>='0'&&c<='9') x=(x<<1)+(x<<3)+(int)(c-'0'),c=getchar();
	x*=f;
}
 
template<typename T,typename I>
void chkmin(T &a,I b){
	a=std::min(a,b);
}
 
template<typename T,typename I>
void chkmax(T &a,I b){
	a=std::max(a,b);
}
 
//随机数部分
std::mt19937 rnd(std::chrono::steady_clock::now().time_since_epoch().count());
template<typename T>
int getgen(T L,T R){
	return rnd()%(R-L+1)+L;
}
 
const double eps=1e-6,pi=std::acos(-1);
const ll inf=1e18,MOD1=998244353,MOD2=1e9+7,MOD3=2147483579,MOD4=19198111;
 
bool Mbe;
 
stack(char) st;
template<typename T>
void print(T x){
	while(st.size()) st.pop();
	if(x==0) st.push('0');
	if(x<0) putchar('-'),x=-x;
	while(x) st.push((char)('0'+x%10)),x/=10;
	while(st.size()) putchar(st.top()),st.pop();
}
template<typename T>
void printsp(T x){
	print(x);
	putchar(' ');
}
template<typename T>
void println(T x){
	print(x);
	putchar('\n');
}
 
const int maxn=2e5+10;
 
std::vector<int>a[maxn];
 
int n,dfn[maxn],fpr[maxn],get[maxn*4],tim,num,fa[maxn],lg[maxn*4],ST[maxn][50],fs[maxn],b[maxn],TS[maxn][50];
 
std::vector<int>c[maxn];
 
namespace LCA{
	void dfs(int p){
		if(!dfn[p]) dfn[p]=++tim;
		fpr[p]=++num,get[num]=tim; 
		fs[dfn[p]]=p;
		for(int v:a[p]){
			if(v==fa[p]) continue;
			fa[v]=p,dfs(v),get[++num]=dfn[p];
		}
	}
	void init(){
		lg[0]=-1;
		for(int i=1;i<=num;i++) lg[i]=lg[i>>1]+1;
		for(int i=1;i<=num;i++) ST[i][0]=get[i];
		for(int j=1;j<=lg[num];j++)
			for(int i=1;i+(1<<j)-1<=num;i++) ST[i][j]=std::min(ST[i][j-1],ST[i+(1<<(j-1))][j-1]);
		return ;
	}
	int query(int x,int y){
		x=fpr[x],y=fpr[y];
		if(x>y) std::swap(x,y);
		int len=y-x+1,_=lg[len];
		return fs[std::min(ST[x][_],ST[y-(1<<_)+1][_])];
	}
}
 
int dep[maxn*2];
 
namespace DS{
	void dfs(int p,int fat){
		dep[p]=dep[fat]+1;
		for(int v:a[p]) if(v!=fat) dfs(v,p);
		return ;
	}
	void init(){
		for(int i=1;i<=n;i++) TS[i][0]=fa[i];
		for(int j=1;j<=lg[n];j++)
			for(int i=1;i<=n;i++){
				TS[i][j]=TS[TS[i][j-1]][j-1];
			}
		return ;
	}
	int query(int x,int y){
		int _=LCA::query(x,y);
		int ans=0;
		for(int j=25;j>=0;j--){
			if(TS[x][j]==0) continue;
			if(dep[TS[x][j]]<dep[_]) continue;
			ans+=(1<<j);
			x=TS[x][j];
		}
		for(int j=25;j>=0;j--){
			if(TS[y][j]==0) continue;
			if(dep[TS[y][j]]<dep[_]) continue;
			ans+=(1<<j);
			y=TS[y][j];
		}
		return ans;
	}
}
 
int xs[maxn*4],len=0,numl[maxn],dis[maxn];
 
bool gz[maxn*2];
 
std::vector<int>d[maxn];
 
int ans=0;
 
namespace ft{
	int id;
	bool cmp(int s1,int s2){
			return dfn[s1]<dfn[s2];
	}
	void build(int t){
		id=-1;
		std::sort(c[t].begin(),c[t].end(),cmp);
		len=0;
		for(int i=0;i<c[t].size()-1;i++){
			xs[++len]=c[t][i];
			xs[++len]=LCA::query(c[t][i],c[t][i+1]);
		}
		xs[++len]=c[t][c[t].size()-1];
		std::sort(xs+1,xs+len+1,cmp);
		len=std::unique(xs+1,xs+len+1)-xs-1;
		for(int i=1;i<len;i++){
			int lc=LCA::query(xs[i],xs[i+1]);
			d[lc].push_back(xs[i+1]);
			d[xs[i+1]].push_back(lc);
		}
		id=xs[1];
		for(int i=2;i<=len;i++) if(dep[xs[i]]<dep[id]) id=xs[i];
//		for(int i=1;i<=len;i++) debug("%lld ",xs[i]);
//		putchar('\n');
	}
	void dfs(int p,int fat){
		numl[p]=1,dis[p]=0;
		if(gz[p]) numl[p]=1;
		else numl[p]=0;
		//get ans of p
		for(int v:d[p]){
			if(v==fat) continue;
			dfs(v,p);
			numl[p]+=numl[v];
			dis[p]+=numl[v]*DS::query(v,p)+dis[v];
		}
		for(int v:d[p]){
			if(v==fat) continue;
			ans+=(numl[v]*DS::query(v,p)+dis[v])*(numl[p]-numl[v]);
//			debug("rt=%lld v=%lld ans=%lld dis=%lld\n",p,v,ans,DS::query(v,p));
		}
	}
}
 
bool Men;
 
signed main(){
//	file();
	assert(1);
	debug("%.8lfMB\n",(&Mbe-&Men)/1048576.0);
	read(n);
	for(int i=1;i<n;i++){
		int u,v;
		read(u),read(v);
		a[u].push_back(v),a[v].push_back(u);
	}
	LCA::dfs(1);
	LCA::init();
	DS::dfs(1,0);
	DS::init();
//	debug("query:%lld\n",DS::query(1,8));
	for(int i=1;i<=n;i++) read(b[i]);
	for(int i=1;i<=n;i++) c[b[i]].push_back(i);
	int tot=0;
	for(int i=1;i<=n;i++){
		if(c[i].size()==0) continue;
		for(int j:c[i]) gz[j]=1;
		ans=0;
		ft::build(i);
		ft::dfs(ft::id,0);
		tot+=ans;
		for(int j:c[i]) gz[j]=0;	
		for(int i=1;i<len;i++){
			int lc=LCA::query(xs[i],xs[i+1]);
			d[lc].clear();
			d[xs[i+1]].clear();
		}
	}
	println(tot);
	debug("%.8lfms\n",1e3*clock()/CLOCKS_PER_SEC);
	return 0;
}
/*
Good luck
-std=c++14 -Wall -O2 -Wextra -Wl,-stack=2147483647
*/

细节较多，容易 RE，注意清空。

时间复杂度 $O (n \log n)$ ，常数较大。

posted @ 2024-06-22 23:55 BYR_KKK 阅读(36) 评论(0) 编辑收藏举报

刷新页面返回顶部

登录后才能查看或发表评论，立即登录或者逛逛博客园首页

相关博文：

· ABC298Ex(2)

· ABC340

· abc359_G Sum of Tree Distance 题解

· [ABC349] AtCoder Beginner Contest 349 题解

· ATcoder 比赛记录

阅读排行：
· 25岁的心里话
· 闲置电脑爆改个人服务器（超详细） #公网映射 #Vmware虚拟网络编辑器
· 零经验选手，Compose 一天开发一款小游戏！
· 因为Apifox不支持离线，我果断选择了Apipost！
· 通过 API 将Deepseek响应流式内容输出到前端

公告

昵称： BYR_KKK
园龄： 1年4个月
粉丝： 22
关注： 38

+加关注

2025年3月

日

一

二

三

四

五

六

byr 的小窝

noip2025 加把劲/qt

ABC359

G

公告

搜索

常用链接

随笔分类

随笔档案

阅读排行榜

评论排行榜

推荐排行榜

最新评论

就你一个人会，就你会！

	#include<bits/stdc++.h>
	#define int long long
	#define fi first
	#define se second

	//#define file(...) freopen("##__VA_ARGS__.in","r",stdin);freopen("##__VA_ARGS__.out","w",stdout)
	//#define outfile(...) freopen("##__VA_ARGS__.out","w",stdout)
	//#define infile(...) freopen("##__VA_ARGS__.in","r",stdin);
	//#define datafile(...) freopen("##__VA_ARGS__.in","w",stdout)
	#define debug(...) fprintf(stderr,##__VA_ARGS__)
	//#define printf(...) fprintf(stdout,##__VA_ARGS)
	template<typename T,typename I>
	#define heap(T) std::priority_queue<T>
	#define umap(T,I) std::unordered_map<T,I>
	#define pii(T,I) std::pair<T,I>
	#define map(T,I) std::map<T,I>
	#define vector(T) std::vector<T>
	#define stack(T) std::stack<T>
	using pii=std::pair<int,int>;
	using ll=long long;
	using i128=__int128;
	using ull=unsigned long long;

	template<typename T>
	inline void read(T &x){
	x=0;
	char c=getchar();
	int f=1;
	while(c<'0'\|\|c>'9'){
	if(c=='-') f=-1;
	c=getchar();
	}
	while(c>='0'&&c<='9') x=(x<<1)+(x<<3)+(int)(c-'0'),c=getchar();
	x*=f;
	}

	template<typename T,typename I>
	void chkmin(T &a,I b){
	a=std::min(a,b);
	}

	template<typename T,typename I>
	void chkmax(T &a,I b){
	a=std::max(a,b);
	}

	//随机数部分
	std::mt19937 rnd(std::chrono::steady_clock::now().time_since_epoch().count());
	template<typename T>
	int getgen(T L,T R){
	return rnd()%(R-L+1)+L;
	}

	const double eps=1e-6,pi=std::acos(-1);
	const ll inf=1e18,MOD1=998244353,MOD2=1e9+7,MOD3=2147483579,MOD4=19198111;

	bool Mbe;

	stack(char) st;
	template<typename T>
	void print(T x){
	while(st.size()) st.pop();
	if(x==0) st.push('0');
	if(x<0) putchar('-'),x=-x;
	while(x) st.push((char)('0'+x%10)),x/=10;
	while(st.size()) putchar(st.top()),st.pop();
	}
	template<typename T>
	void printsp(T x){
	print(x);
	putchar(' ');
	}
	template<typename T>
	void println(T x){
	print(x);
	putchar('\n');
	}

	const int maxn=2e5+10;

	std::vector<int>a[maxn];

	int n,dfn[maxn],fpr[maxn],get[maxn4],tim,num,fa[maxn],lg[maxn4],ST[maxn][50],fs[maxn],b[maxn],TS[maxn][50];

	std::vector<int>c[maxn];

	namespace LCA{
	void dfs(int p){
	if(!dfn[p]) dfn[p]=++tim;
	fpr[p]=++num,get[num]=tim;
	fs[dfn[p]]=p;
	for(int v:a[p]){
	if(v==fa[p]) continue;
	fa[v]=p,dfs(v),get[++num]=dfn[p];
	}
	}
	void init(){
	lg[0]=-1;
	for(int i=1;i<=num;i++) lg[i]=lg[i>>1]+1;
	for(int i=1;i<=num;i++) ST[i][0]=get[i];
	for(int j=1;j<=lg[num];j++)
	for(int i=1;i+(1<<j)-1<=num;i++) ST[i][j]=std::min(ST[i][j-1],ST[i+(1<<(j-1))][j-1]);
	return ;
	}
	int query(int x,int y){
	x=fpr[x],y=fpr[y];
	if(x>y) std::swap(x,y);
	int len=y-x+1,_=lg[len];
	return fs[std::min(ST[x][_],ST[y-(1<<_)+1][_])];
	}
	}

	int dep[maxn*2];

	namespace DS{
	void dfs(int p,int fat){
	dep[p]=dep[fat]+1;
	for(int v:a[p]) if(v!=fat) dfs(v,p);
	return ;
	}
	void init(){
	for(int i=1;i<=n;i++) TS[i][0]=fa[i];
	for(int j=1;j<=lg[n];j++)
	for(int i=1;i<=n;i++){
	TS[i][j]=TS[TS[i][j-1]][j-1];
	}
	return ;
	}
	int query(int x,int y){
	int _=LCA::query(x,y);
	int ans=0;
	for(int j=25;j>=0;j--){
	if(TS[x][j]==0) continue;
	if(dep[TS[x][j]]<dep[_]) continue;
	ans+=(1<<j);
	x=TS[x][j];
	}
	for(int j=25;j>=0;j--){
	if(TS[y][j]==0) continue;
	if(dep[TS[y][j]]<dep[_]) continue;
	ans+=(1<<j);
	y=TS[y][j];
	}
	return ans;
	}
	}

	int xs[maxn*4],len=0,numl[maxn],dis[maxn];

	bool gz[maxn*2];

	std::vector<int>d[maxn];

	int ans=0;

	namespace ft{
	int id;
	bool cmp(int s1,int s2){
	return dfn[s1]<dfn[s2];
	}
	void build(int t){
	id=-1;
	std::sort(c[t].begin(),c[t].end(),cmp);
	len=0;
	for(int i=0;i<c[t].size()-1;i++){
	xs[++len]=c[t][i];
	xs[++len]=LCA::query(c[t][i],c[t][i+1]);
	}
	xs[++len]=c[t][c[t].size()-1];
	std::sort(xs+1,xs+len+1,cmp);
	len=std::unique(xs+1,xs+len+1)-xs-1;
	for(int i=1;i<len;i++){
	int lc=LCA::query(xs[i],xs[i+1]);
	d[lc].push_back(xs[i+1]);
	d[xs[i+1]].push_back(lc);
	}
	id=xs[1];
	for(int i=2;i<=len;i++) if(dep[xs[i]]<dep[id]) id=xs[i];
	// for(int i=1;i<=len;i++) debug("%lld ",xs[i]);
	// putchar('\n');
	}
	void dfs(int p,int fat){
	numl[p]=1,dis[p]=0;
	if(gz[p]) numl[p]=1;
	else numl[p]=0;
	//get ans of p
	for(int v:d[p]){
	if(v==fat) continue;
	dfs(v,p);
	numl[p]+=numl[v];
	dis[p]+=numl[v]*DS::query(v,p)+dis[v];
	}
	for(int v:d[p]){
	if(v==fat) continue;
	ans+=(numl[v]DS::query(v,p)+dis[v])(numl[p]-numl[v]);
	// debug("rt=%lld v=%lld ans=%lld dis=%lld\n",p,v,ans,DS::query(v,p));
	}
	}
	}

	bool Men;

	signed main(){
	// file();
	assert(1);
	debug("%.8lfMB\n",(&Mbe-&Men)/1048576.0);
	read(n);
	for(int i=1;i<n;i++){
	int u,v;
	read(u),read(v);
	a[u].push_back(v),a[v].push_back(u);
	}
	LCA::dfs(1);
	LCA::init();
	DS::dfs(1,0);
	DS::init();
	// debug("query:%lld\n",DS::query(1,8));
	for(int i=1;i<=n;i++) read(b[i]);
	for(int i=1;i<=n;i++) c[b[i]].push_back(i);
	int tot=0;
	for(int i=1;i<=n;i++){
	if(c[i].size()==0) continue;
	for(int j:c[i]) gz[j]=1;
	ans=0;
	ft::build(i);
	ft::dfs(ft::id,0);
	tot+=ans;
	for(int j:c[i]) gz[j]=0;
	for(int i=1;i<len;i++){
	int lc=LCA::query(xs[i],xs[i+1]);
	d[lc].clear();
	d[xs[i+1]].clear();
	}
	}
	println(tot);
	debug("%.8lfms\n",1e3*clock()/CLOCKS_PER_SEC);
	return 0;
	}
	/*
	Good luck
	-std=c++14 -Wall -O2 -Wextra -Wl,-stack=2147483647
	*/