HDOJ树形DP专题之Godfather

题目链接

这题Balance Act那题差不多，不过这题的数据量更大，时间有点卡，我的O(N)的算法都跑了1100多MS（时限2S）。

#include <stdio.h>
#include <string.h>
#include <vector> 
#define N 50005
#define MAX(a,b) ((a)>(b)?(a):(b))
using namespace std;
vector<int> dep[N];
int u[2*N],v[2*N],first[N],next[2*N];
int n,p[N],d[N],sum[N],w[N],ans[N],cnt,dmax;
void dfs(int i,int fa)
{
  int e,j;
  d[i]=(fa==-1?0:d[fa]+1);
  dep[d[i]].push_back(i);
  dmax=MAX(dmax,d[i]);
  for(e=first[i];e>=0;e=next[e])
  {
    j=v[e];
    if(j!=fa) dfs(j,p[j]=i);
  }
}
void dp()
{
  int i,j,k,l;
  memset(w,0,sizeof(w));
  for(i=0;i<n;i++)  sum[i]=1;
  for(i=dmax;i>=0;i--)
  {
    l=dep[i].size();
    for(k=0;k<l;k++)
    {
      j=dep[i][k];
      w[j]=MAX(w[j],n-sum[j]);
      if(i>0) w[p[j]]=MAX(w[p[j]],sum[j]),sum[p[j]]+=sum[j];
    }
  }
}
int main()
{
  int i,e,a,b,min;
  while(~scanf("%d",&n))
  {
    memset(first,-1,sizeof(first));
    memset(next,-1,sizeof(next));
    for(e=0;e<n-1;e++)
    {
      scanf("%d%d",&a,&b);
      a--,b--;
      u[2*e]=a,v[2*e]=b;
      next[2*e]=first[u[2*e]];
      first[u[2*e]]=2*e;
      u[2*e+1]=b,v[2*e+1]=a;
      next[2*e+1]=first[u[2*e+1]];
      first[u[2*e+1]]=2*e+1;
    }
    dmax=0;
    for(i=0;i<n;i++)  dep[i].clear();
    dfs(0,-1);
    dp();
    min=0x7fffffff;
    cnt=0;
    for(i=0;i<n;i++)
    {
      if(w[i]<min)  min=w[i],cnt=0,ans[cnt++]=i;
      else if(w[i]==min)  ans[cnt++]=i;
    }
    printf("%d",ans[0]+1);
    for(i=1;i<cnt;i++)  printf(" %d",ans[i]+1);
    printf("\n");
  }
  return 0;
}