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

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;
}