O(n)树上dp代码求调教
查看原帖
O(n)树上dp代码求调教
1145602
dongzhenmao楼主2024/8/4 18:45

复杂度没算错,快读也用了,不知道错在哪,悬关

#include <stdio.h>
#include <ctype.h>
#include <algorithm>
#include <string.h>
#include <vector>
#define lnt long long
#define inf 0x3f3f3f3f
using namespace std;
int xx;char chh,ff;int read(){
	xx=ff=0;while(!isdigit(chh)){if(chh=='-'){ff=1;}chh=getchar();}
	while(isdigit(chh)){xx=(xx<<1)+(xx<<3)+chh-'0';chh=getchar();}return ff? -xx: xx;
}
const int N=1e6,MOD=1e9+7;
lnt n,ans[N],anst,siz[N],lj[N];//
vector<lnt> G[N];
void dfs1(int,int);
void dfs2(int,int);
void mod(lnt &x){x%=MOD;}
int main(){
	n=read();
	for(int i=1;i<n;++i){
		int x=read(),y=read();
		G[x].push_back(y);
		G[y].push_back(x);
	}
	dfs1(1,0);
	dfs2(1,0);
	printf("%lld",(anst+MOD)%MOD);
	
	return 0;
}
void dfs1(int now,int fa){
	siz[now]=1;
	for(auto to: G[now]){
		if(to==fa){continue;}
		dfs1(to,now);
		siz[now]+=siz[to];
		ans[now]+=ans[to]+lj[to];
		lj[now]+=lj[to];
		mod(lj[now]);mod(ans[now]);
	}
	for(int i=0;i<G[now].size();++i){
		int to=G[now][i];
		for(int e=i+1;e<G[now].size();++e){
			int to2=G[now][e];
			if(to==to2){continue;}
			if(to==fa || to2==fa){continue;}
			lj[now]+=siz[to]*siz[to2]%MOD;
			lj[now]%=MOD;
		}
	}
	lj[now]+=siz[now]-1;
	lj[now]%=MOD;
}
void dfs2(int now,int fa){
	anst+=ans[now];
	mod(anst);
	for(auto to: G[now]){
		if(to==fa){continue;}
		lnt A=lj[now]-lj[to]-siz[to]*(siz[now]-siz[to]-1)-siz[to];
		mod(A);
		ans[to]+=ans[now]-ans[to]-lj[to]+A;
		mod(ans[to]);
		lj[to]+=(siz[now]-siz[to])*(siz[to]-1)+A+siz[now]-siz[to];
		mod(lj[to]);
		siz[to]=n;
		dfs2(to,now);
	}
}






2024/8/4 18:45
加载中...