[Hackerrank题目选做] Tree Pruning

题目链接:https://www.hackerrank.com/challenges/tree-pruning

题解: 先吐槽一下数据，非常无脑的O(nk^2)的代码加了一点小优化就0.12s闪过，惊讶.jpg

         

 先上O(nk^2)的AC代码:

#include <bits/stdc++.h>
#define ll long long
using namespace std;
vector <int> vec[100005];
int lar[100005];
ll pre[201];
ll f[100005][201],a[100005];
int n,k;
void dfs(int pos,int fa)
{int i,j,l;
f[pos][0]=a[pos];lar[pos]=0;
for (i=0;i<vec[pos].size();i++)
{if (vec[pos][i]!=fa)
{dfs(vec[pos][i],pos);
for (j=0;j<=lar[pos];j++)
{pre[j]=f[pos][j];
f[pos][j]=-1000000000000000ll;
}
for (j=0;j<=lar[pos];j++)
{for (l=0;(l<=lar[vec[pos][i]]&&j+l<=k);l++)
{f[pos][j+l]=max(f[pos][j+l],pre[j]+f[vec[pos][i]][l]);}
}
lar[pos]+=lar[vec[pos][i]];
if (lar[pos]>k) {lar[pos]=k;}
}
}
if (f[pos][1]<0) {f[pos][1]=0;}
lar[pos]=max(lar[pos],1);
}
int main (){
	int i,j,u,v;
	memset (f,-0x7f,sizeof(f));
	scanf ("%d%d",&n,&k);
	for (i=1;i<=n;i++)
	{scanf ("%lld",&a[i]);}
	for (i=1;i<n;i++)
	{scanf ("%d%d",&u,&v);
	vec[u].push_back(v);
	vec[v].push_back(u);
	}
	dfs(1,0);
	ll ans=0;
	for (i=0;i<=k;i++)
	{if (f[1][i]>ans)
	{ans=f[1][i];}
	}
	printf ("%lld",ans);
	return 0;
}
	

然后，正解也不是很难想，

先O(n)弄出dfs序，然后每个子树就是一段连续区间了，

只需要对于每个子树是否选择进行dp转移就可以了，

时间复杂度O(nk).

Code:

#include <bits/stdc++.h>
#define ll long long
using namespace std;
vector <int> vec[100005];
ll f[100005][205],a[100005];
int st[100005],p[100005],ed[100005];
int n,k,cnt=0;
void dfs(int pos,int fa)
{st[pos]=++cnt;
p[cnt]=pos;
for (int i=0;i<vec[pos].size();i++)
{if (vec[pos][i]!=fa) {dfs(vec[pos][i],pos);}
}
ed[pos]=cnt;
}
int main (){
	int i,j,u,v;
	scanf ("%d%d",&n,&k);
	for (i=1;i<=n;i++)
	{scanf ("%lld",&a[i]);}
	for (i=1;i<n;i++)
	{scanf ("%d%d",&u,&v);
	vec[u].push_back(v);
	vec[v].push_back(u);
	}
	dfs(1,0);
	memset (f,-0x7f,sizeof(f));
	f[0][0]=0;
	for (i=1;i<=n;i++)
	{for (j=0;j<=k;j++)
	{f[i][j]=max(f[i][j],f[i-1][j]+a[p[i]]);
	f[ed[p[i]]][j+1]=max(f[ed[p[i]]][j+1],f[i-1][j]);
	}
	}
	for (j=0;j<=k;j++)
	{f[n][k]=max(f[n][j],f[n][k]);}
	printf ("%lld",f[n][k]);
	return 0;
}