快手第四题 好序列
直接统计不好序列 路径全为红即为不好序列
每个点记录子节点里有几个点通过不好路径相连
考虑每棵子树对答案的贡献即可
#include<bits/stdc++.h> using namespace std; const int mod = 1e9+7; struct node { int to; int c; node(int _to,int _c):to(_to),c(_c){} }; struct node2 { int to; int c; int num; node2(int _to,int _c,int _num):to(_to),c(_c),num(_num){} }; vector<node>g[100010]; vector<node2>G[100010]; bool vis[100010]; void build(int root){ if(vis[root]==1) return; vis[root]=1; for(int i=0;i<g[root].size();i++){ if(vis[g[root][i].to]==0){ G[root].push_back(node2(g[root][i].to,g[root][i].c,0)); build(g[root][i].to); } } } int num[100010],ans=0; int fast(int a,int b) { int ans=1; while(b) { if(b%2==1) ans=1LL*ans*a%mod; b/=2; a=1LL*a*a%mod; } return ans; } int n,k,l,r,c; void solve(int root) { num[root]=1; for(int i=0;i<G[root].size();i++){ solve(G[root][i].to); if(G[root][i].c==0) { num[root]+=num[G[root][i].to]; ans=(ans+mod-fast(num[G[root][i].to],k))%mod; } } ans=(ans+fast(num[root],k))%mod; } int main() { memset(vis,0,sizeof(vis)); memset(num,0,sizeof(num)); cin>>n>>k; for(int i=0;i<n-1;i++){ scanf("%d %d %d",&l,&r,&c); g[l].push_back(node(r,c)); g[r].push_back(node(l,c)); } build(1); solve(1); cout<<(fast(n,k)+mod-ans)%mod<<endl; }