ZOJ - 4045 District Division(树)

Sample Input
3
4 2
1 3
3 2
1 4
6 3
1 3
1 4
1 6
2 5
5 1
8 4
1 2
2 3
2 4
1 5
5 6
5 7
5 8
Sample Output
YES
1 4
2 3
NO
YES
4 3 2 1
5 6 7 8

题意：给一棵有n个节点的树，问是否能将树分成几个区域，每个区域k个节点，如果可以输出YES并输出每个区域的节点都是谁，否则输出NO

首先dfs计算每个节点子树的大小，然后直接遍历整棵树，判断其子树是否能整除k，能整除说明其子树可以被分为n/k个区域，否则不行，不行的情况可能是该节点被包含在某个区域内了，之后对比能整除的节点个数是否等于n/k的区域数量，如果相等说明可以分配，否则有多余节点导致一些区域不够凑k个节点。

如果可行，则深搜遍历之前记录的节点，输出即是区域内的节点名。

#include<stdio.h>
#include<string.h>
#include<algorithm>
#include<vector>
using namespace std;
const int maxn=100005;
int n,t,k;
vector<int>mp[maxn];
vector<int>ans;
bool vis[maxn];
int sum[maxn],fa[maxn];
void dfs1(int rt,int pre)
{
    sum[rt]=1;
    fa[rt]=pre;
    for(int i=0; i<mp[rt].size(); i++)
    {
        int v=mp[rt][i];
        if(mp[rt][i]!=pre)
        {
            dfs1(v,rt);
            sum[rt]+=sum[v];
        }
    }
}
void dfs2(int rt,int pre)
{
    if(sum[rt]%k==0)ans.push_back(rt),vis[rt]=true;
    for(int i=0; i<mp[rt].size(); i++)
        if(mp[rt][i]!=pre)
            dfs2(mp[rt][i],rt);
}
int dfs3(int rt,int pre)
{
    for(int i=0; i<mp[rt].size(); i++)
    {
        if(!vis[mp[rt][i]]&&mp[rt][i]!=pre)
        {
            printf(" %d",mp[rt][i]);
            dfs3(mp[rt][i],rt);
        }
    }
}
bool judge()
{
    if(ans.size()!=n/k)return false;
    return true;
}
int main()
{
    scanf("%d",&t);
    while(t--)
    {
        int from,to;
        scanf("%d%d",&n,&k);
        ans.clear();
        for(int i=0; i<=n; i++)mp[i].clear();
        memset(vis,false,sizeof vis);
        for(int i=0; i<n-1; i++)
        {
            scanf("%d%d",&from,&to);
            mp[from].push_back(to);
            mp[to].push_back(from);
        }
        dfs1(1,0);
        dfs2(1,0);
        bool ju=judge();
        printf("%s\n",ju?"YES":"NO");
        if(ju)
        {
            for(int i=0; i<ans.size(); i++)
            {
                printf("%d",ans[i]);
                dfs3(ans[i],fa[ans[i]]);
                printf("\n");
            }
        }
    }
}