【树上计数乘法原理】ABC129 E - Virus Tree 2

àì夳堔傛蜴生んèń 2024-03-16 15:21 8阅读 0赞

**明天就要篮球被了，这不会是我今年最后一场比赛了吧**

[E - Sum Equals Xor (atcoder.jp)][E - Sum Equals Xor _atcoder.jp]

题意：

给定一棵树和K种颜色，让你给这棵树涂色，使得两点之间距离<=2的两个点不同色

![a5f7d35cdf7246a48441cb30e4318d89.png][]

思路：

一开始很自然地想到了树形计数DP，但是很不幸，N,K<=1e5，爆空间了，因此考虑组合计数

第二个想法是容斥，先算出所有涂色方案是K^N，然后减去存在d=1的同色的，再减去d=2的同色的，然后加上d=1且d=2的同色的，然后发现这好像更难算了

然后就不会了，事实上，模拟一下第二个样例就会发现，这个就是乘法原理而已

**一开始一定要先去模拟样例再去想思路！这样会省下很多时间**

手摸一下第二个样例：

![205dfba298e54a9bbbe618818597e675.jpeg][]

可以发现，深度为1的方案数是K，深度为2的方案数是K-1，其他的深度的最左边的点都是K-2，然后方案数随着从左到右递减

所以可以直接乘法原理了

Code：

#include <bits/stdc++.h>
    
    #define int long long
    
    using namespace std;
    
    const int mxn=1e5+10;
    const int mxe=1e5+10;
    const int mod=1e9+7;
    
    int N,K,u,v,tot=0;
    int head[mxn];
    int dep[mxn],a[mxn];
    
    struct ty{
    	int to,next;
    }edge[mxe<<2];
    
    void add(int u,int v){
    	edge[tot].to=v;
    	edge[tot].next=head[u];
    	head[u]=tot++;
    }
    void G_init(){
    	tot=0;
    	for(int i=0;i<=N;i++){
    		head[i]=-1;
    	}
    }
    void dfs1(int u,int fa){
    	dep[u]=dep[fa]+1;
    	if(dep[u]==1) a[u]=K;
    	else if(dep[u]==2) a[u]=K-1;
    	else a[u]=K-2;
    	for(int i=head[u];~i;i=edge[i].next){
    		if(edge[i].to==fa) continue;
    		dfs1(edge[i].to,u);
    	}
    }
    void dfs2(int u,int fa){
    	int ok=0,last;
    	for(int i=head[u];~i;i=edge[i].next){
    		if(edge[i].to==fa) continue;
    		if(!ok){
    			ok=1;
    			last=a[edge[i].to];
    		}else{
    			a[edge[i].to]=last-1;
    			last=a[edge[i].to];
    		}
    	}
    	for(int i=head[u];~i;i=edge[i].next){
    		if(edge[i].to==fa) continue;
    		dfs2(edge[i].to,u);
    	}
    }
    void solve(){
    	cin>>N>>K;
    	G_init();
    	for(int i=1;i<=N-1;i++){
    		cin>>u>>v;
    		add(u,v);
    		add(v,u);
    	}
    	dfs1(1,0);
    	dfs2(1,0);
    	//for(int i=1;i<=N;i++) cout<<a[i]<<" \n"[i==N];
    	int ans=1ll;
    	for(int i=1;i<=N;i++){
    		int u=max(0ll,a[i]);
    		ans=(ans*u)%mod;
    	}
    	cout<<ans<<'\n';
    }
    signed main(){
    	ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    	int __=1;//cin>>__;
    	while(__--)solve();return 0;
    }

[E - Sum Equals Xor _atcoder.jp]: https://atcoder.jp/contests/abc129/tasks/abc129_e
[a5f7d35cdf7246a48441cb30e4318d89.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/16/4c2c0232a79c42c5ba688b5b65018a54.png
[205dfba298e54a9bbbe618818597e675.jpeg]: https://image.dandelioncloud.cn/pgy_files/images/2024/03/16/507f813c95ff45cd99c90c2910419bb3.png