题意描述

对一棵树，给出每条edge的颜色（红或黑），从树中取三个节点，定义若三个节点每两两之间的路径中都有红色的edge，则该取法为一种有效取法，求总共有多少种取法。

解析

本题的做法是对于，每对用黑色相连的节点用并查集维护，形成多组仅用黑色相连节点集合，在任意三组中各取一个节点即为有效取法，求和即可。 此题中一个值得注意的点是，若直接用三重循环枚举，或是稍微优化用一个前缀和数组，复杂度都大于 O(n^2),所以我们从反面入手，求可能性总和，并排除掉三个点在有两个或三个在同一集合的情况，复杂度降到O(n)。

#include <cmath>
#include <cstdio>
#include <vector>
#include <iostream>
#include <algorithm>
using namespace std;

int fa[100100];
int sz[100100];
long long sum[100010];
int find(int x){
    return fa[x]==x?x:fa[x]=find(fa[x]);//where mistake happens 
}
void merge(int x,int y){
    int px=find(x),py=find(y);
    if(px==py) return;
    sz[px]+=sz[py];
    fa[py]=px;
}
long long ans=0;
const long long mod=1e9+7;
int main() {
    vector<long long> v;
    for(int i=0;i<100010;i++) fa[i]=i,sz[i]=1;
    int NodeNum;cin>>NodeNum;
    for(int i=0;i<NodeNum-1;i++){
        int a,b;char c;
        cin>>a>>b>>c;
        if(c=='b'){//若为黑色则合并
            merge(a,b);
        }
    }
  
    for(int i=1;i<=NodeNum;i++){
        if(fa[i]==i) v.push_back(sz[i]),ans+=sz[i];//求每个集合节点数量
    }
  
    ans=((ans*(ans-1)%mod)*(ans-2)/6)%mod;//combinational number C(3,ans)
    
    for(int i=0;i<v.size();i++) {//choose triplets(a,b,c) from one sets
        ans-=((v[i]*(v[i]-1)*(v[i]-2))/6)%mod;
        if(ans<0) ans+=mod;
    }
  
    for(int i=0;i<v.size();i++) {//choose triplets(a,b,c) from two sets
        ans-=((v[i]*(v[i]-1)/2)*(NodeNum-v[i]))%mod;
        if(ans<0) ans+=mod;
    }
    cout<<ans<<endl;
    return 0;
}