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;
}