HYSBZ-4033-树上染色(树上DP)

链接:

https://vjudge.net/problem/HYSBZ-4033

题意:

有一棵点数为N的树，树边有边权。给你一个在0~N之内的正整数K，你要在这棵树中选择K个点，将其染成黑色，并
将其他的N-K个点染成白色。将所有点染色后，你会获得黑点两两之间的距离加上白点两两之间距离的和的收益。
问收益最大值是多少。

思路:

数上任意两点的距离,每条边e(u,v)的贡献是cntl[u], cntr[v], 就是u左边的点乘上v右边的点乘上权值.
同时令DP[u][k]为u的子树中有k个黑点,的最大值.
在DFS中从u到v可推出Dp[u][i+j] = max(Dp[u][i+j], Dp[u][j]+Dp[v][i]+dis*cnt)其中Dp[u][j]为以u为根,别的边的贡献.
Dp[v][i]表示此条边连着的v的子数的贡献,最后算出这条边的贡献.
cnt就是两端黑白的组合数.

代码:

#include <iostream>
#include <cstdio>
#include <cstring>
#include <vector>
//#include <memory.h>
#include <queue>
#include <set>
#include <map>
#include <algorithm>
#include <math.h>
#include <stack>
#include <string>
#include <assert.h>
#include <iomanip>
#define MINF 0x3f3f3f3f
using namespace std;
typedef long long LL;
const int MAXN = 2e3+10;

struct Node
{
    int to;
    LL dis;
};
vector<Node> G[MAXN];
LL Dp[MAXN][MAXN];
int Cnt[MAXN];
int n, k;

void Dfs(int u, int v)
{
    Cnt[v] = 1;
    for (int i = 0;i < G[v].size();i++)
    {
        int node = G[v][i].to;
        if (node == u)
            continue;
        Dfs(v, node);
        for (int j = min(Cnt[v], k);j >= 0;j--)
        {
            for (int z = min(Cnt[node], k);z >= 0;z--)
            {
                LL ti = 1LL*z*(k-z)+1LL*(Cnt[node]-z)*(n-Cnt[node]-(k-z));
                Dp[v][j+z] = max(Dp[v][j+z], Dp[v][j]+Dp[node][z]+G[v][i].dis*ti);
            }
        }
        Cnt[v] += Cnt[node];
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    int u, v;
    LL w;
    for (int i = 1;i < n;i++)
    {
        cin >> u >> v >> w;
        G[u].push_back(Node{v, w});
        G[v].push_back(Node{u, w});
    }
    Dfs(0, 1);
    cout << Dp[1][k] << endl;

    return 0;
}