BZOJ4033或洛谷3177 [HAOI2015]树上染色

BZOJ原题链接

洛谷原题链接

很明显的树形(DP)。
因为记录每个点的贡献很难，所以我们可以统计每条边的贡献。
对于每一条边，设边一侧的黑点有(B_x)个，白点有(W_x)，另一侧黑点有(B_y)，白点有(W_y)，边权为(w)，那么这条边的贡献就是((W_x imes W_y + B_x imes B_y) imes w)。
然后设计(DP)状态，定义(f[x][v])，表示以(x)为根的子树里分配(v)个黑点的最大贡献。
初始化为(-1)，在(dfs)到(x)点时，再初始化(f[x][0] = f[x][1] = 0)。
设(y)表示(x)的一个儿子， (k)为题目中所述。
于是有转移方程：

(qquadqquad i = min{k, size[x]} longrightarrow 0)

(qquadqquadquad j = 0longrightarrow min{i, size[y]})

(qquadqquadqquad f[x][i] = max{f[x][i], f[x][i - j] + f[y][j] + value}qquad ifquad f[x][i - j] e -1)

(i)是在以(x)为根的子树中分配多少黑点，(j)是在以(y)为根的子树中分配多少黑点。
(value)是通过(x o y)这条边所新增的贡献，即(value = (j imes (k - j) + (size[y] - j) imes (n - size[y] - k + j)) imes w_{x o y})。

#include<cstdio>
#include<cstring>
using namespace std;
typedef long long ll;
const int N = 2010;
int fi[N], di[N << 1], ne[N << 1], da[N << 1], si[N], l, k, n;
ll f[N][N];
inline int re()
{
	int x = 0;
	char c = getchar();
	bool p = 0;
	for (; c < '0' || c > '9'; c = getchar())
		p |= c == '-';
	for (; c >= '0' && c <= '9'; c = getchar())
		x = x * 10 + c - '0';
	return p ? -x : x;
}
inline void add(int x, int y, int z)
{
	di[++l] = y;
	da[l] = z;
	ne[l] = fi[x];
	fi[x] = l;
}
inline ll maxn(ll x, ll y)
{
	return x > y ? x : y;
}
inline int minn(int x, int y)
{
	return x < y ? x : y;
}
void dfs(int x, int fa)
{
	int i, j, v, y, o;
	si[x] = 1;
	f[x][0] = f[x][1] = 0;
	for (i = fi[x]; i; i = ne[i])
		if ((y = di[i]) ^ fa)
		{
			dfs(y, x);
			si[x] += si[y];
		}
	for (v = fi[x]; v; v = ne[v])
		if ((y = di[v]) ^ fa)
			for (i = minn(k, si[x]); ~i; i--)
				for (j = 0, o = minn(i, si[y]); j <= o; j++)
					if (~f[x][i - j])
						f[x][i] = maxn(f[x][i], f[x][i - j] + f[y][j] + (1LL * j * (k - j) + 1LL * (si[y] - j) * (n - si[y] - k + j)) * da[v]);
}
int main()
{
	int i, x, y, z;
	n = re();
	k = re();
	for (i = 1; i < n; i++)
	{
		x = re();
		y = re();
		z = re();
		add(x, y, z);
		add(y, x, z);
	}
	memset(f, -1, sizeof(f));
	dfs(1, 0);
	printf("%lld", f[1][k]);
	return 0;
}