BZOJ 4016: [FJOI2014]最短路径树问题

建树+点分治 难度在于建树

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<queue>
#define sc second
#define pr pair<int,int>
#define mp make_pair
using namespace std;
int root,Cnt,F_top,cnt,n,m,K,N,instack[1000005],last[1000005],vis[1000005],dis[1000005],F_[1000005],F[1000005],G[1000005],F_stack[1000005],G_stack[1000005],F_sum[1000005],G_sum[1000005],sz[1000005];
long long ans,ans_sum;
priority_queue<pr,vector<pr>,greater<pr> > q;
struct node{
	int to,next,val;
}e[1000005];
void add(int a,int b,int c){
	e[++cnt].to=b;
	e[cnt].next=last[a];
	e[cnt].val=c;
	last[a]=cnt;
}
void Dijkstra(int S){
	for (int i=1; i<=n; i++) vis[i]=0,dis[i]=1e9;
	dis[S]=0;
	q.push(mp(0,S));
	while (!q.empty()){
		int x=q.top().sc;
		q.pop();
		if (vis[x]) continue;
		vis[x]=1;
		for (int i=last[x]; i; i=e[i].next){
			int V=e[i].to;
			if (dis[V]>dis[x]+e[i].val){
				dis[V]=dis[x]+e[i].val;
				q.push(mp(dis[V],V));
			}
		}
	}
}
struct node1{
	int x,val;
};
vector<node1> vec[1000005];
bool cmp(node1 a,node1 b){
	return a.x<b.x;
}
struct node2{
	int x,y,val;
}E[1000005];
void dfs(int x){
	vis[x]=1;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		vec[x].push_back((node1){V,e[i].val});
	}
	sort(vec[x].begin(),vec[x].end(),cmp);
	for (int i=0; i<(int)vec[x].size(); i++){
		int V=vec[x][i].x;
		if (dis[V]==dis[x]+vec[x][i].val && !vis[V]) {
			E[++Cnt]=(node2){x,V,vec[x][i].val};
			dfs(V);
		}
	}
}
void find_root(int x,int fa){
	sz[x]=1,F_[x]=0;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V] || V==fa) continue;
		find_root(V,x);
		sz[x]+=sz[V];
		F_[x]=max(F_[x],sz[V]);
	}
	F_[x]=max(F_[x],N-sz[x]);
	if (F_[x]<F_[root]) root=x;
}
void get_dis(int x,int fa,int val,int dep){
	if (dep==K){
		int ANS=val,ANS_sum=1;
		if (ANS==ans) ans_sum+=ANS_sum;
		else if (ANS>ans) ans=ANS,ans_sum=ANS_sum;
	}
	if (!instack[dep]) F_stack[++F_top]=dep,instack[dep]=1;
	if (F[dep]==val) F_sum[dep]++;
	else if (F[dep]<val) F[dep]=val,F_sum[dep]=1;
	sz[x]=1;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V] || V==fa) continue;
		get_dis(V,x,val+e[i].val,dep+1);
		sz[x]+=sz[V];
	}
}
void solve(int x){
	int G_top=0;
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V]) continue;
		F_top=0;
		get_dis(V,x,e[i].val,2);
		for (int j=1; j<=F_top; j++) 
			if (K+1-F_stack[j]>=1){
				int ANS=F[F_stack[j]]+G[K+1-F_stack[j]];
				long long ANS_sum=1ll*F_sum[F_stack[j]]*G_sum[K+1-F_stack[j]];
				if (ANS==ans) ans_sum+=ANS_sum;
				else if (ANS>ans) ans=ANS,ans_sum=ANS_sum;
			}
		for (int j=1; j<=F_top; j++)
			if (F[F_stack[j]]==G[F_stack[j]]) G_sum[F_stack[j]]+=F_sum[F_stack[j]];
			else if (F[F_stack[j]]>G[F_stack[j]]) G[F_stack[j]]=F[F_stack[j]],G_sum[F_stack[j]]=F_sum[F_stack[j]];
		for (int j=1; j<=F_top; j++) F[F_stack[j]]=0,instack[F_stack[j]]=0,G_stack[++G_top]=F_stack[j];
	}
	for (int i=1; i<=G_top; i++) G[G_stack[i]]=0;
}
void divide(int x){
	vis[x]=1;
	solve(x);
	for (int i=last[x]; i; i=e[i].next){
		int V=e[i].to;
		if (vis[V]) continue;
		N=sz[V],root=0;
		find_root(V,x);
		divide(root);
	}
}
int main(){
	scanf("%d%d%d",&n,&m,&K);
	for (int i=1; i<=m; i++){
		int x,y,z;
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,z);
		add(y,x,z);
	}
	Dijkstra(1);
	memset(vis,0,sizeof(vis));
	dfs(1);
	F_[0]=1e9;
	cnt=0;
	memset(last,0,sizeof(last));
	for (int i=1; i<=Cnt; i++) add(E[i].x,E[i].y,E[i].val),add(E[i].y,E[i].x,E[i].val);
	memset(vis,0,sizeof(vis));
	N=n,root=0;
	find_root(1,0);
	divide(root);
	printf("%lld %lld\n",ans,ans_sum);
	return 0;
}
/*
4 6 2
1 2 5
1 4 2
2 4 6
2 3 6
3 4 5
1 3 1
*/