题意简述
一棵树，其中有若干条关键边，求有多少点三元组 \((i,j,k)\) 满足 \(i\) 与 \(j\) 间有关键边且 \(i\) 与 \(k\) 间有关键边。
\(n \leq 10^5\)。
题目分析
同学说这题很有迷惑性。
考虑转化 \(i\) 与 \(j\) 间有关键边的条件，很容易想到将所有关键边断开后会形成若干联通块 \(t_1,t_2,\cdots,t_k\)。
设点 \(u\) 所属的联通块为 \(s_u\)，即 \(u \in s_u\)。
显然 \(u\) 和 \(v\) 间有关键边当且仅当 \(s_u \neq s_v\)。
\[\begin{aligned} ans & =\sum_{i \neq j \neq k} [s_i \neq s_j][s_i \neq s_k]\\ & =\sum_{i} \sum_{j \neq k}[s_i \neq s_j][s_i \neq s_k]\\ & =\sum_{i} 2\dbinom{n-|s_i|}{2}\\ & =2\sum_{i_1}^n \dbinom{n-|s_i|}{2}\\ & =2\sum_{i_1}^k |t_i| \dbinom{n-|t_i|}{2}\\ \end{aligned} \]dfs 计算联通块大小，组合数直接计算即可。
时间复杂度 \(O(n)\)。
拓展 1：点带权。
记 \(W=\sum\limits_{i=1}^n w_i\)，\(W_k=\sum\limits_{i \in s_k}^n w_i\)（即该点所在联通块的点权和），\(W_k=\sum\limits_{i \in s_k}^n w_i^2\)。
\[\begin{aligned} ans & =\sum\limits_{i}w_i(\sum\limits_{j,k}[s_i \neq s_j][s_i \neq s_k]w_jw_k-\sum\limits_{j} [s_i \neq s_j]w_j^2)\\ & =\sum\limits_i w_i[(W-W_i)^2-W_i']\\ \end{aligned} \]拓展 2：多元组。
若将条件改为 \(p\) 元组，且任意两点间有关键边，即为任意两点不同属一个联通块。
即从 \(k\) 个数里任意选出 \(p\) 个数的积之和。
这是一个经典问题，即 \([x^p]\prod_{i=1}^k (1+t_ix)\)，分治+FFT 即可在 \(O(k\log^2 k)\) 的时间复杂度内解决。
代码
#include<bits/stdc++.h>
using namespace std;
const long long N=1e5+5;
struct nod{
long long to,nxt,w;
}e[N*2];
long long head[N],cnt;
void add(long long u,long long v,long long w){
e[++cnt].to=v;
e[cnt].w=w;
e[cnt].nxt=head[u];
head[u]=cnt;
}
long long check(long long x){
while(x){
if(x%10!=4 && x%10!=7) return 0;
x/=10;
}
return 1;
}
long long n,id,ans;
bool vis[N];
void dfs(long long u){
vis[u]=1;
id++;
for(long long i=head[u];i;i=e[i].nxt){
long long v=e[i].to;
if(vis[v] || e[i].w) continue;
dfs(v);
}
}
int main(){
cin>>n;
for(long long i=1;i<n;i++){
long long u,v,w;
scanf("%lld%lld%lld",&u,&v,&w);
long long t=check(w);
add(u,v,t);
add(v,u,t);
}
for(long long i=1;i<=n;i++){
if(!vis[i]){
id=0;
dfs(i);
ans+=id*(n-id)*(n-id-1);
}
}
printf("%lld",ans);
}