poj2378 Tree Cutting

/*
 * poj2378 AC
 *  一次水過。
 *  簡單的樹狀DP，先做一次dfs算出每個結點的子孫結點總數，包括自己，sum[i]。
 *  再做dfs計算刪去每個結點x是否滿足條件，即滿足：
 *   1) 每個結點的所有子結點i的sum[i]是否小於half。
 *   2) 父結點方向上的結點總數，sum[1]-sum[x]是否小於half(1爲樹的根結點)。
 *  若滿足則爲一個解，同時計算每個子結點是否滿足。
 *  若某結點不滿足條件2，那麼其子結點一定不滿足條件，
 * 因爲sum[1]-sum[x]只會遞增，而使子結點一定不滿足條件2。
 *
 * */
#include<cstdio>
#include<algorithm>
#include<memory.h>
#define N 10010
using namespace std;

int tree[N<<1],head[N<<1],next[N<<1],sum[N+5];
int ans[N],len = 0,half;
bool vis[N];

long dfs_count(int x)
{
    int i;
    sum[x] = 1,vis[x] = true;
    for(i=head[x];i;i=next[i])
       if(!vis[tree[i]]) sum[x] += dfs_count(tree[i]);
    return sum[x];
}

bool dfs_dp(int x)
{
    vis[x] = true;
    if(sum[1]-sum[x]>half) return false; 
    int i,flag = true;
    for(i=head[x];i;i=next[i])
        if(!vis[tree[i]])
        {
           if(dfs_dp(tree[i])) ans[++len] = tree[i];
           if(sum[tree[i]]>half) flag = false; 
        }
    return flag;
} 

int main()
{
    int n,tot,i,j,k;
//    FILE* fin;
//    fin = fopen("d.in","r");
    scanf("%d",&n);
//    fscanf(fin,"%d",&n);
    for(half=n>>1,tot=0,i=1;i<=n-1;i++) {
//        fscanf(fin,"%d%d",&j,&k);
        scanf("%d%d",&j,&k);
        tree[++tot] = k,next[tot] = head[j],head[j] = tot;
        tree[++tot] = j,next[tot] = head[k],head[k] = tot;
    }
    memset(sum,0,sizeof(sum));
    memset(vis,false,sizeof(vis));
    dfs_count(1);

    memset(vis,false,sizeof(vis));
    if(dfs_dp(1)) ans[++len] = 1;
    if(len==0) 
    {
        printf("NONE\n");
        return 0;
    }

    sort(ans+1,ans+len+1);
    for(i=1;i<=len;i++)
        printf("%d\n",ans[i]);
//    fclose(fin);
    return 0;
}