题目描述
难度分:1900
输入n(1≤n≤105)和一棵树的n−1条边(节点编号从1开始),每条边包含3个数x、y、z(1≤z≤109),表示有一条边权为z的边连接x和y。
定义幸运数为只包含4和7的数,例如47、744、4。
输出有多少个三元组(i,j,k),满足i,j,k互不相同,且从i到j的简单路径上存在幸运边权,从i到k的简单路径上也存在幸运边权。
注意(1,2,3)和(2,1,3)是两个不同的三元组。
输入样例1
4
1 2 4
3 1 2
1 4 7
输出样例1
16
输入样例2
4
1 2 4
1 3 47
1 4 7447
输出样例2
24
算法
并查集
这个题最开始想的是树形DP
,但最终没想出来有什么比较好写代码的状态,还是放弃了。
首先想到可以枚举i节点,j和k两个节点的方案数通过组合数学的方式得到(即便用树形DP
,大思路也肯定是这样)。而边权对我们而言是没用的,我们只需要知道这条边是不是幸运边就好。然后就发现了一个很有用的性质,把所有的幸运边都断开,整棵树就会变成很多的连通块,i,j,k三个点从这些连通块里面选必然是满足条件的。
因此,读入边集的时候我们就可以判断边是不是幸运边,如果不是就连接两个节点,同时用并查集维护连通性。接下来遍历并查集的所有根节点i,得到每个连通块的大小sz[i]。那么三元组的方案数就是2Σisz[i]×C2n−sz[i],C2n−sz[i]表示剩下两个节点j和k从其他连通块的n−sz[i]个节点中选择,最后还要乘以2是因为j和k交换顺序也是两种不同的方案。
复杂度分析
时间复杂度
建图需要遍历n−1条边,对于每条边,要用O(log10U)的时间复杂度来判断其是不是幸运边,可能还需要利用并查集合并两个非幸运边连接的节点,时间复杂度为O(logn)。总体的时间复杂度为O(n(log10U+logn)),其中U是边权最大值。接下来遍历并查集的所有根节点计算答案即可,每个根节点的计算都是O(1)的,在最差情况下会有O(n)级别的根节点,因此时间复杂度为O(n)。
综上,整个算法的时间复杂度为O(n(log10U+logn)+n)。
空间复杂度
并查集父节点数组p、连通块大小数组sz,以及图的邻接表空间都是线性的,为O(n),这也是整个算法的额外空间复杂度。
C++ 代码
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 100010;
int n, p[N], sz[N];
int find(int x) {
if(p[x] != x) p[x] = find(p[x]);
return p[x];
}
void merge(int x, int y) {
int rx = find(x), ry = find(y);
if(rx != ry) {
p[rx] = ry;
sz[ry] += sz[rx];
}
}
LL C2(LL x) {
return x*(x - 1)/2;
}
bool check(int w) {
unordered_set<int> st;
while(w) {
st.insert(w % 10);
w /= 10;
}
if(st.size() == 1) {
return st.count(4) || st.count(7);
}else if(st.size() == 2) {
return st.count(4) && st.count(7);
}
return false;
}
int main() {
scanf("%d", &n);
for(int i = 1; i <= n; i++) {
p[i] = i;
sz[i] = 1;
}
for(int i = 1; i < n; i++) {
int u, v, w;
scanf("%d%d%d", &u, &v, &w);
if(!check(w)) {
merge(u, v);
}
}
unordered_set<int> roots;
for(int i = 1; i <= n; i++) {
roots.insert(find(i));
}
LL ans = 0;
for(int root: roots) {
LL cnt = sz[root];
ans += C2(n - cnt)*cnt*2LL;
}
printf("%lld\n", ans);
return 0;
}