CF1681F Unique Occurrences

• 题意：一棵树，问每条路径上只出现一次的值的个数的和。
• 思路：
  
  显然想到考虑边贡献。每条边权下放到下面的哪个点。\(up_i\)为上面第一个点权等于它的点。我们需要一个子树内点权等于它的点（如果满足祖孙关系，不要孙）（除它自己的）sz和。
  
  这样每个点的\(sz\)向\(up\)贡献。
  
  这样差分求出上面的点和下面的点方案数。乘一下就好了
• code:

点击查看代码

#include<bits/stdc++.h>
using namespace std;
const int N=1e6+5;
typedef long long ll;
int up[N],n,sz[N],fa[N],lst[N],sum[N],s1[N],nxt[N],to[N],head[N],len[N],ecnt;
void add_edge(int u,int v,int w) {
	nxt[++ecnt]=head[u];to[ecnt]=v;len[ecnt]=w;head[u]=ecnt;
	nxt[++ecnt]=head[v];to[ecnt]=u;len[ecnt]=w;head[v]=ecnt;
}
void init(int u) {
	sz[u]=1;
	for(int i=head[u];i;i=nxt[i]) {
		int v=to[i];if(v==fa[u])continue;
		fa[v]=u;init(v);sz[u]+=sz[v];
	}
}
void init_(int u) {
	for(int i=head[u];i;i=nxt[i]) {
		int v=to[i],w=len[i],pre=lst[w];if(v==fa[u])continue;
		if(pre)sum[pre]+=sz[v];
		else s1[w]+=sz[v];
		up[v]=pre;lst[w]=v;
		init_(v);
		lst[w]=pre;
	}
}

ll ans;
void init__(int u) {
	for(int i=head[u];i;i=nxt[i]) {
		int v=to[i],w=len[i];if(v==fa[u])continue;
		int y=up[v];
		ll w1,w2=sz[v]-sum[v];
		if(!y) {w1=(sz[1]-sz[v])-(s1[w]-sz[v]);}
		else {w1=(sz[y]-sz[v])-(sum[y]-sz[v]);}
//		printf("(%d-%d) w1=%lld w2=%lld\n",u,v,w1,w2);
		ans+=w1*w2;
		init__(v);
	}
}
int main() {
	scanf("%d",&n);
	for(int i=1;i<n;i++) {int u,v,w;scanf("%d%d%d",&u,&v,&w);add_edge(u,v,w);}
	init(1),init_(1),init__(1);
//	for(int i=1;i<=n;i++)printf("i=%d sum=%lld s1=%lld\n",i,sum[i],s1[i]);
	printf("%lld",ans);
	return 0;
}