AtCoder Beginner Contest 259 F - Select Edges(树形dp)

$link$

题意

现在给出$n$个结点的树形结构，并给出$n-1$条带权边。要求从这些边中选出若干条，使得他们权值之和最大。要求是，这些边中出现的结点次数不能超过题目中给出的限制。

输入第二行是结点的限制，剩下的$n-1$行是无向边和带权值。

思路

这个题不难看出来是树形dp。

不难想到，负权边一定不能选择。

假设当前结点的限制是k条边，那么我们尽可能选择连接这个点前k大的正权边，由于树形结构的dfs是自上而下回溯的，因此回溯时，该结点和其子节点是没有边相连接的，现在问题就转化成为了是否选择该结点和子结点的边，

设$dp[i][1]$为不选择当前结点和父结点的边即子结点与其子结点已经完成全部边的选择，$dp[i][0]$为选择当前结点和父结点的边，即子结点和其子结点完成了全部边-1的选择。

因此当前结点的$dp[u][0]$就是选择$d[u]-1$条边与之相连，所产生的价值

因此当前结点的$dp[u][1]$就是选择$d[u]$条边与之相连，所产生的价值

最后汇总到根节点就是找到$dp[1][1]$和$dp[1][0]$的最大值。

#include <bits/stdc++.h>
using namespace std;
#define int long long
vector<int> vec[300005];
int n;
int d[310210];
int e[606060], ne[606060], h[303030], w[606060], idx;
int dp[303030][2];
int ans = 0;
void add(int a, int b, int c)
{
    e[idx] = b, ne[idx] = h[a], w[idx] = c, h[a] = idx++;
}
void dfs(int u, int fa)
{
    int tot = 0;
    vector<int> vec;
    for (int i = h[u]; ~i; i = ne[i])
    {
        int j = e[i];
        if (j == fa)
            continue;
        dfs(j, u);
        if (dp[j][1] > 0)
            tot += dp[j][1];
        if (d[j] > 0 && dp[j][0] + w[i] > dp[j][1] && dp[j][0] + w[i] > 0)
        {
            vec.push_back(dp[j][0] + w[i] - dp[j][1]);
        }
    }
    sort(vec.rbegin(), vec.rend());
    int sum = 0;
    dp[u][1] = dp[u][0] = tot;
    for (int i = 0; i < d[u]; i++)
    {
        if (vec.size() > i)
            sum += vec[i];

        if (i == d[u] - 1)
        {
            dp[u][1] += sum;
        }
        if (i == d[u] - 2)
        {
            dp[u][0] += sum;
        }
    }
}
signed main()
{
    cin >> n;
    memset(h, -1, sizeof h);
    for (int i = 1; i <= n; i++)
    {
        cin >> d[i];
        vec[i].resize(d[i] + 1);
    }
    for (int i = 1; i < n; i++)
    {
        int a, b, c;
        cin >> a >> b >> c;
        add(a, b, c);
        add(b, a, c);
    }
    dfs(1, 0);
    cout << max(dp[1][0], dp[1][1]) << endl;
}