[BZOJ3697]采药人的路径(点分治+树形dp)

[BZOJ3697]采药人的路径(点分治+树形dp)

题面

采药人的药田是一个树状结构，每条路径上都种植着同种药材。

采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。

采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

分析

显然可以点分治。设阴性和阳性的边权分别为1和-1,那么平衡就是路径长度为0

考虑点分治的时候dfs每一颗分治子树,用\(mark_i\)来标记当前节点到分治中心的路径上距离的出现情况,如果\(mark_{deep_x}=1\),就说明\(x\)与某个祖先到中心出现过一个相同的距离,也就是说这一段上一定存在休息点。

然后考虑合并分治中心出来的若干子树。设\(g_{0/1,i}\)表示当前子树中子树中距离为i的路径数目,0/1表示是否有休息点。合并的时候类似树形dp,为了防止一个子树内重复合并，要再定义一个数组\(f_{0/1,i}\)表示当前已经合并过子树中距离为i的路径数目.每新来一个子树,先dfs一遍算出g,再把g和f合并。
每棵子树对答案的贡献是

\[\begin{aligned}&g_{0,0} \times (f_{0,0}-1) \text{(两条路径都平衡,分治中心是休息站,-1是去掉路径经过点数为0的情况)} \\ &+ \sum_{i}g_{1,-i} \times f_{1,i}+ g_{0,-i}\times f_{1,i}+ g_{1,-i}\times f_{0,i} \text{两条路径和为0,平衡,且至少一条路径上存在休息站}\end{aligned} \]

实现上要注意负下标的问题

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define maxn 200000 
using namespace std;
typedef long long ll;
int n;
struct edge{
	int from;
	int to;
	int len;//两种边,距离分别为-1和1 
	int next;
}E[maxn*2+5];
int esz=1;
int head[maxn*2+5];
void add_edge(int u,int v,int w){
	esz++;
	E[esz].from=u;
	E[esz].to=v;
	E[esz].len=w;
	E[esz].next=head[u];
	head[u]=esz;
}
bool vis[maxn+5];
int sz[maxn+5];
int ms[maxn+5];
void dfs_root(int x,int fa,int tot_sz,int &root){
	sz[x]=1;
	ms[x]=0;
	for(int i=head[x];i;i=E[i].next){
		int y=E[i].to;
		if(y!=fa&&!vis[y]){
			dfs_root(y,x,tot_sz,root);
			sz[x]+=sz[y];
			ms[x]=max(ms[x],sz[y]);
		}
	}
	ms[x]=max(ms[x],tot_sz-sz[x]);
	if(ms[x]<ms[root]) root=x;
} 
int get_root(int x,int tot_sz){
	int root=0;
	ms[0]=tot_sz+1;
	dfs_root(x,0,tot_sz,root);
	return root;
} 

struct marray{//支持负数下标的数组 
	ll cnt[maxn*2+5];
	inline ll & operator [] (int id){
		return cnt[id+maxn];
	} 
};
marray f[2];//当前已经合并的子树中距离为i的路径数目,0/1表示是否有休息点 
marray g[2];//当前子树中子树中距离为i的路径数目
marray mark;//标记某个距离是否出现 
int maxd=0;
void get_dist(int x,int fa,int dist){
	maxd=max(maxd,abs(dist));
	if(mark[dist]) g[1][dist]++;//祖先中有距离为dist的点,那么当前节点与那个祖先之间必有一个休息点 
	else g[0][dist]++;
	mark[dist]++;
	for(int i=head[x];i;i=E[i].next){
		int y=E[i].to;
		if(y!=fa&&!vis[y]){
			get_dist(y,x,dist+E[i].len);
		}
	}
	mark[dist]--; 
}
ll ans=0;
void solve(int x){
	vis[x]=1;
	f[0][0]=1;
	int mmaxd=0;
	for(int p=head[x];p;p=E[p].next){
		int y=E[p].to;
		if(!vis[y]){
			maxd=0;
			get_dist(y,x,E[p].len);
//			mmind=min(mind,mmind);
			mmaxd=max(maxd,mmaxd);
			ans+=g[0][0]*(f[0][0]-1);//两边距离为0,说明两条路径都平衡,休息点在x 
			//-1是去掉路径经过点数为0的情况,即一开始的f[0][0]=1 
			for(int i=-mmaxd;i<=maxd;i++) ans+=g[1][-i]*f[1][i]+g[0][-i]*f[1][i]+g[1][-i]*f[0][i];
			//两边距离总和为0,整条路径平衡,且至少一条路径上有休息点 
			for(int i=-maxd;i<=maxd;i++){//合并子树 
				f[0][i]+=g[0][i];
				f[1][i]+=g[1][i];
				g[0][i]=g[1][i]=0;
			}
		}
	}
	for(int i=-mmaxd;i<=mmaxd;i++) f[0][i]=f[1][i]=0;
	for(int i=head[x];i;i=E[i].next){
		int y=E[i].to;
		if(!vis[y]) solve(get_root(y,sz[y]));
	}
}
int main(){
	int u,v,w;
	scanf("%d",&n);
	for(int i=1;i<n;i++){
		scanf("%d %d %d",&u,&v,&w);
		if(w==1){
			add_edge(u,v,1);
			add_edge(v,u,1);
		}else{
			add_edge(u,v,-1);
			add_edge(v,u,-1);
		}
	}
	solve(get_root(1,n)); 
	printf("%lld\n",ans);
}