洛谷P3915 树的分解

题目描述

给出 $n$ 个点的树和 $k$，问能否把树划分成 $\frac{n}{k}$ 个连通块，且每个连通块的点数都是 $k$。

输入格式

第 $1$ 行，$1$ 个整数 $T$，表示数据组数。接下来 $T$ 组数据，对于每组数据：
第 $1$ 行，$2$ 个整数 $n$, $k$；
接下来 $n-1$ 行，每行 $2$ 个整数 $a_i,b_i$，表示边 $a_i,b_i$。点用 $1,2,\cdots,n$ 编号。

输出格式

对于每组数据，输出YES或NO。

输入输出样例

输入 #1

2
4 2
1 2
2 3
3 4
4 2
1 2
1 3
1 4

输出 #1

YES
NO

说明/提示

对于 $60\%$ 的数据，$1 \le n,k \le 10^3$；
对于 $100\%$ 的数据，$1 \le T \le 10$, $1 \le n ,k \le 10^5$。

思路

摘取自洛谷用户Youngsc的题解
如果存在能将一棵树分成若干个大小为 $k$ 的联通块的方案数，那么每一个联通块都有一定是一棵子树，且分割方案是唯一的，因此我们队这棵树进行深搜，并同时记录子树大小，如果当前大小恰好为 $k$ 的话，就将该子树大小记录为$0$，相当于剪去这一部分。最后只要判断剪掉的部分时候等于 $\frac{n}{k}$ 就可以了。

代码

/******************************************************************
Copyright: 11D_Beyonder All Rights Reserved
Author: 11D_Beyonder
Problem ID: 洛谷 P3915
Date: 2021 Feb. 1st
Description: Tree, DFS
*******************************************************************/
#include<iostream>
#include<cstring>
#include<vector>
using namespace std;
const int N=100003;
int _;
int n,k;
vector<int>G[N];
int cnt,sz[N];
void DFS(int,int);
void init();
int main(){
	for(cin>>_;_;_--){
		cin>>n>>k;
		init();
		for(int i=1;i<=n-1;i++){
			int u,v;
			cin>>u>>v;
			G[u].push_back(v);
			G[v].push_back(u);
		}
		DFS(1,0);
		puts(cnt==n/k&&!(n%k)?"YES":"NO");
	}
	return 0;
}
void init(){
	//初始化
	for(int i=1;i<=n;i++){
		G[i].clear();
		sz[i]=0;
	}
	cnt=0;
}
void DFS(int x,int f){
	sz[x]=1;
	for(auto y:G[x]){
		if(y==f){
			continue;
		}
		DFS(y,x);
		sz[x]+=sz[y];
	}
	if(sz[x]==k){
		cnt++;
		sz[x]=0;//删去子树
	}
}