AtCoder Beginner Contest 133 E - Virus Tree 2(组合数学)

​

题意：

给出一棵树和种颜色，给树的节点涂色，要求距离的节点不能同色，求涂色方案。

思路：

• 考虑递推每个节点的涂色方案，根据乘法原理求出答案
• 根节点的涂色方案为，没有限制
• 假设当前遍历到为根的子树，子节点分三种情况
• 子节点是的第一个儿子并且没有父节点，这样的涂色方案为，去除父节点的颜色即可
• 子节点是的第一个儿子并且有父节点，这样的涂色方案为，去除父节点的颜色和的父节点的颜色
• 子节点不是的第一个儿子，答案为上一个子节点的答案，因为两个子节点之间的距离为，颜色也不能相同

代码：

const int maxn=3e5+7,mod=1e9+7,inf=0x3f3f3f3f;

int n,k,dp[maxn];
vector<int>g[maxn];

void dfs(int u,int fa){
  int las=0;
  for(auto t:g[u]){
    if(t==fa) continue;
    if(las) dp[t]=las-1;
    else if(!fa) dp[t]=k-1;
    else dp[t]=k-2;
    las=dp[t];
  }
  for(auto t:g[u]){
    if(t!=fa) dfs(t,u);
  }
}

int main(){
  n=read,k=read;
  rep(i,1,n-1){
    int u=read,v=read;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dp[1]=k;
  dfs(1,0);
  ll ans=1;
  rep(i,1,n) ans=ans*dp[i]%mod;
  write(ans);
  return 0;
}