JZOJ 4814. 【NOIP2016提高A组五校联考2】tree

Problem

Description

给一棵n 个结点的有根树，结点由1 到n 标号，根结点的标号为1。每个结点上有一个物品，第i 个结点上的物品价值为vi。
你需要从所有结点中选出若干个结点，使得对于任意一个被选中的结点，其到根的路径上所有的点都被选中，并且选中结点的个数不能超过给定的上限lim。在此前提下，你需要最大化选中结点上物品的价值之和。
求这个最大的价值之和。

Input

第一行为两个整数n; lim
接下来n 行，第i 行包含一个整数vi，表示结点i 上物品的价值。
接下来n- 1 行，每行包含两个整数u; v, 描述一条连接u; v 结点的树边。

Output

输出一行答案。

Sample Input

6 4
-5
4
-6
6
9
6
3 2
3 1
2 4
2 5
1 6

Sample Output

2

Data Constraint

对于前20% 的数据，1<=n; lim<=10
对于前60% 的数据，1<=n; lim<=100
对于100% 的数据，1<=n; lim<=3000; |vi| <=10^5 数据有梯度，保证给出的是合法的树。

Solution

树形DP，$f[i][j]$表示DFS遍历到i，在DFS过程中选择了j个节点，获得的最大收益。DFS遍历到一个节点x时，先将他父亲的f值赋给x，然后再DP。
DP方程显然。

$$f[i][j]=max(f[i][j],f[son][j-1]+v[son]) (son是i的儿子)$$

Code

#include<iostream>
#include<cstdio>
#include<cmath>
#include<cstring>
#define N 3010
#define fo(i,a,b) for(i=a;i<=b;i++)
using namespace std;
int v[N],a[N][N],b[N][N],f[N][N],deep[N];
int mx,i,j,n,lim,x,y,tot;
bool bz[N];
void dg(int k,int fa)
{
    int i;
    fo(i,1,a[k][0])
    {
        if (a[k][i]==fa) continue;
        deep[a[k][i]]=deep[k]+1;
        dg(a[k][i],k);
    }
}
void dfs(int k)
{
    int i,j;
    bz[k]=1;
    fo(i,1,b[k][0])
    {
        if (bz[b[k][i]]) continue;
        memcpy(f[b[k][i]],f[k],sizeof(f[k]));
        dfs(b[k][i]);
        fo(j,1,lim)
            f[k][j]=max(f[k][j],f[b[k][i]][j-1]+v[b[k][i]]);
    }
}
int main()
{
    scanf("%d%d",&n,&lim);
    fo(i,1,n) scanf("%d",&v[i]);
    fo(i,1,n-1)
    {
        scanf("%d%d",&x,&y);
        a[x][++a[x][0]]=y;
        a[y][++a[y][0]]=x;
    }
    dg(1,0);
    fo(i,1,n)
        fo(j,1,a[i][0])
            if (deep[i]<deep[a[i][j]]) b[i][++b[i][0]]=a[i][j];
    dfs(1);
    printf("%d",v[1]+f[1][lim-1]);
}