ABC133 E Virus Tree 2 树形遍历+染色

题目：E - Virus Tree 2

给定一个含个节点的无向树，给每个顶点赋上一整数（染上一种颜色）。并且，若两个顶点距离，那么两顶点值不同（颜色不同）。

考察

首先无向树的根可以随意设定，我们就从第1个开始查找（邻接表）。根的颜色肯定是K种可用。我们固定好根的颜色后，用DFS向下探索遍历所有的枝及叶。赋值的时候，我们可以考虑以下两点

• 根的颜色数为
• 固定顶点值和父级节点值，其所有子节点的赋值方案为，其中为子节点数量

代码实现

// some macroes and includes
//---------------------
#define MAXN 100005
#define MOD 1000000007
//---------------------

ll n,k;
vector<int> G[MAXN];
ll ret=1;

ll modmult(ll a , ll b ){
  return ( (a%MOD) * (b%MOD)  ) %MOD;
  }

void solve(ll lastId,ll curId,ll depth){
    
    if(ret==0) {ret=0;return;}
    ll temp=depth==0 ? k:1;
    ll id=k-min(depth+1,ll(2));
    REP(i,G[curId].size()){
        if(G[curId][i] == lastId) continue;
        solve( curId , G[curId][i] , depth+1);
        temp = modmult(temp,max(id--,ll(0)));
    }
    ret = modmult(ret,temp);
    //cout<<ret<<endl;
}

int main(){
    cin>>n>>k;
    REP(i,n-1) {int a,b;cin>>a>>b;G[a].PB(b);G[b].PB(a);}
    solve(-1,1,0);
    cout << ret <<endl;
    return 0;
}