https://codeforces.com/contest/1139/problem/C

題意

一顆有n個點的樹,需要挑選出k個點組成序列(可重復),按照序列的順序遍歷樹,假如經過黑色的邊,那麽這個序列就是好的,問有多少個好的序列

題解

• 黑邊不連,紅邊連,假如兩個點不在同一並查集,那麽一定經過黑邊
• 定義\(dp[i][j][k]\)為選擇前i個點,起始點為j,是否已經經過黑邊(k)的方案數
• \(dp[i-1][j][0]*(n-N[fin(j)])+dp[i-1][j][1]*n - > dp[i][j][1]\)
• \(dp[i-1][j][0]*N[fin(j)] - > dp[i][j][0]\)

代碼

#include<bits/stdc++.h>
#define MOD 1000000007
using namespace std;
int fa[100005],i,j,n,m,u,v,w;
int fin(int u){return fa[u]==u?u:fa[u]=fin(fa[u]);}
void merge(int u,int v){
    int x=fin(u),y=fin(v);
    if(x!=y)fa[x]=y;
}
long long dp[100005][3],ans,N[100005];
int main(){
    cin>>n>>m;
    for(i=1;i<=n;i++)fa[i]=i;
    for(i=0;i<n-1;i++){
        scanf("%d%d%d",&u,&v,&w);
        if(!w)merge(u,v);
    }
    for(i=1;i<=n;i++){N[fin(i)]++;dp[i][0]=1;}
    for(i=2;i<=m;i++){
        for(j=1;j<=n;j++){
            dp[j][1]=(dp[j][0]*((n-N[fin(j)]+MOD)%MOD)%MOD+dp[j][1]*n%MOD)%MOD;
            dp[j][0]=(dp[j][0]*N[fin(j)])%MOD;
        }
    }
    for(i=1;i<=n;i++)
        ans=(ans+dp[i][1])%MOD;
    cout<<ans;
}
 

Codeforces Round #548 (Div. 2) C dp or 排列組合