【CF109C】Lucky Tree【并查集】【dfs】

题目大意：

题目链接：https://codeforces.com/problemset/problem/109/C
给出一棵树，定义树上一条边的边权若只由4和7两个数字构成，则这一条边为幸运边。求这棵树上有多少三元组$(x,y,z)$满足$x$到$y$的路径上和$y$到$z$的路径上都至少有一条幸运边。

---

思路：

最多绿题难度。结果洛谷给评了一个紫题。
考虑对于每一个点$x$，如果我们求出有$sum$个点与$x$之间的路径至少需要经过一条幸运边，那么答案其实就是$sum(sum-1)$，因为这$sum$个点中，任意两个点都可以与$x$组成一个三元组。
所以我们记录这棵树的每一条边是否为幸运边，然后$dfs$，求出每一个不包含幸运边的连通块的大小。并查集维护每一个点出现在哪一个连通块内。
最后枚举每一个点，不与这个点在同一个连通块内的点就必然在到$x$的路径中有至少一条幸运边。那么就可以直接求答案了。
时间复杂度$O(n)$。没算并查集的复杂度。

---

代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;

const int N=100010;
int n,tot,size[N],father[N],head[N];
bool vis[N];
ll ans;

struct edge
{
	int next,to;
	bool flag;
}e[N*2];

void add(int from,int to,bool flag)
{
	e[++tot].to=to;
	e[tot].flag=flag;
	e[tot].next=head[from];
	head[from]=tot;
}

bool check(int x)
{
	for (;x;x/=10)
		if (x%10!=4 && x%10!=7) return 0;
	return 1;
}

int find(int x)
{
	return x==father[x]?x:father[x]=find(father[x]);
}

void dfs(int x,int fa)
{
	vis[x]=1;
	father[x]=find(fa);
	size[father[x]]++;
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa && !e[i].flag) dfs(v,x);
	}
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d",&n);
	for (int i=1,x,y,z;i<n;i++)
	{
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,check(z)); add(y,x,check(z));
	}
	for (int i=1;i<=n;i++)
		father[i]=i;
	for (int i=1;i<=n;i++)
		if (!vis[i]) dfs(i,i);
	for (int i=1;i<=n;i++)
		ans+=(ll)(n-size[father[i]])*(n-size[father[i]]-1);
	printf("%I64d",ans);
	return 0;
}