CodeForces - 109C Lucky Tree

原题链接：

https://vjudge.net/problem/CodeForces-109C

AC代码：

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <list>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <stdio.h>
#include <stdlib.h>

using namespace std;

typedef long long ll;
const int64_t mod = 1e9 + 7;
const int maxn = 1e6 + 10;
const int INF = 0x3f3f3f3f;
int n, fa[maxn], num[maxn];

int find(int x) {
    if (fa[x] == x)return x;
    return fa[x] = find(fa[x]);
}

void uni(int u, int v) {
    int xx = find(u);
    int yy = find(v);
    if (xx == yy) {
        return;
    }
    fa[yy] = xx;
    num[xx] += num[yy];
}

int check(int w) {
    while (w) {
        if (w % 10 != 4 && w % 10 != 7)return 0;
        w /= 10;
    }
    return 1;
}

int main() {
    cin >> n;
    for (int i = 1; i <= n; i++) {
        fa[i] = i;
        num[i] = 1;
    }
    int u, v, w;
    for (int i = 1; i < n; i++) {
        cin >> u >> v >> w;
        if (!check(w)) {
            uni(u, v);
        }
    }
    ll ans = 0;
    for (int i = 1; i <= n; i++) {
        int t = n - num[find(i)];
        ans += (ll) t * (t - 1);
    }
    printf("%I64d\n", ans);
    return 0;
}

对于每个顶点 i 计算其余各点到其有幸运边路径的点的数量num[i]，利用并查集，以非幸运边建树，则每个连通块之间由幸运边相连，各连通块内都为非幸运边，则对于每个i，ans=ΣC(num[i],2)(再选两个)