【DFS算方案数】AtCoder Beginner Contest 133 E - Virus Tree 2

 

https://atcoder.jp/contests/abc133/tasks/abc133_e

E - Virus Tree 2

---

Time Limit: 2 sec / Memory Limit: 1024 MB

Score : 500500 points

Problem Statement

You are given a tree with NN vertices and N−1 edges. The vertices are numbered 1 to N, and the i-th edge connects Vertex ai and bi.

You have coloring materials of K colors. For each vertex in the tree, you will choose one of the K colors to paint it, so that the following condition is satisfied:

• If the distance between two different vertices x and y is less than or equal to two, x and y have different colors.

How many ways are there to paint the tree? Find the count modulo 1 000 000 007

 

What is tree？

What is distance?

 

Constraints

• 1≤N,K≤1e5
• 1≤ai,bi≤N
• The given graph is a tree.

---

Input

Input is given from Standard Input in the following format:

N K
a1 b1
a2 b2
..
..
..
aN−1 bN−1

Output

Print the number of ways to paint the tree, modulo 1 000 000 007

---

Sample Input 1 

4 3
1 2
2 3
3 4

Sample Output 1 

6

There are six ways to paint the tree.

---

Sample Input 2 

5 4
1 2
1 3
1 4
4 5

Sample Output 2 

48

---

Sample Input 3 

16 22
12 1
3 1
4 16
7 12
6 2
2 15
5 16
14 16
10 11
3 10
3 13
8 6
16 8
9 12
4 3

Sample Output 3 

271414432

题意：给你一棵树，k种颜色，要求两个节点的路径<=2的节点不能是相同颜色的，问你染色的方案数

解法：将1作为节点拎起来，跑dfs，1这个节点可以有k种颜色，第二层节点可以有A(k-1,sz)种方案，第三层及以上的节点可以有A(k-2，sz)种方案。

我们跑一遍dfs，求答案即可

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const ll mod=1e9+7;
const int maxn=1e5+10;
ll fac[maxn],inv[maxn];
vector <int> G[maxn];
struct node
{
    int sz,cs;
}a[maxn];

ll quickpow(ll a,ll b)
{
    ll ret=1;
    while(b)
    {
        if(b&1) ret=(ret*a)%mod;
        a=(a*a)%mod;
        b>>=1;
    }
    return ret;
}

void init(int p)
{
    fac[0]=1;
    for(int i=1;i<=p;i++)
        fac[i]=(fac[i-1]*i)%mod;
    inv[p]=quickpow(fac[p],mod-2);
    for(int i=p-1;i>=0;i--)
    {
        inv[i]=(inv[i+1]*(i+1))%mod;
    }
}

ll A(ll n,ll m)
{
    //!!!!颜色可能不够用
    if(m>n)
        return 0;
    return fac[n]%mod*inv[n-m]%mod;
}

void dfs(int x,int fa,int cs)
{
    a[x].cs=cs;
    for(int i=0;i<G[x].size();i++)
    {
        int v=G[x][i];
        if(v==fa) continue;
        dfs(v,x,cs+1);
        a[x].sz++;
    }
}

int main()
{
    init(100001);
    int n,u,v;
    ll k;
    scanf("%d%lld",&n,&k);
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }

    dfs(1,0,1);
    ll ans=k;
    //cout<<ans<<endl;
    for(int i=1;i<=n;i++)
    {
        if(a[i].cs==1)
        {
            ans=ans*A(k-1,a[i].sz)%mod;
           // cout<<A(k-1,a[i].sz)<<endl;
        }
        else
        {
            ans=ans*A(k-2,a[i].sz)%mod;
            //cout<<A(k-2,a[i].sz)<<endl;
        }
    }
    ans%=mod;
    printf("%lld\n",ans);
    return 0;
}