北化ACM集训队每周程序设计 7 问题 E: 病毒树

• 题意：给你一颗\(n\)个节点的树，有\(k\)个节点，给每个节点涂色，要求每三个相邻的节点颜色不同，求方案数。

• 题解：我们dfs这颗树，当遍历到某个节点时，要给这个点涂色，那么这个点不能和他的父亲，以及他左边的同父亲的兄弟，以及父亲的父亲同一个颜色，然后每个点能涂的颜色相乘就是方案数

• 代码：

  #include <bits/stdc++.h>
  using namespace std;
  #define PII pair<int,int>
  #define fi first
  #define se second
  #define pb push_back
  #define ll long long
  #define ull unsigned long long
  #define PLL pair<ll,ll>
  const int N=1000010;
  const int mod=1e9+7;
   
  int n;
  ll k;
  vector<int> edge[N];
  int sz[N];
  ll ans=1;
   
  void dfs(int u,int fa,int dep){
      ll now=k;
      if(dep==2) now--;
      else if(dep>=2) now-=2;
      now-=(sz[fa]-1);
      ans=ans*max(1ll*0,now)%mod;
      for(auto to:edge[u]){
          if(to==fa) continue;
          sz[u]++;
          dfs(to,u,dep+1); 
      }
  }
   
  int main(){
      scanf("%d %lld",&n,&k);
      for(int i=1;i<n;++i){
          int u,v;
          scanf("%d %d",&u,&v);
          edge[u].pb(v);
          edge[v].pb(u);
      } 
      sz[0]=1;
      dfs(1,0,1);
      printf("%lld\n",ans);
      return 0;
  }