【树形DP】【P3177】[HAOI2015] 树上染色

Description

给定一棵 \(n\) 个点的带权树，要求选 \(k\) 个点染成黑色，剩下染成白色，最大化两两同色点之间的距离和。

Limitations

\(0 \leq k \leq n \leq 2000\)

Solution

首先看一个trick：

考虑如下遍历一棵树的伪代码：

func dfs(u):
  size[u] <- 1
  for v in clild[u] do
    dfs(v)
    for i = 1 : size[u] do
      for j = 1: size[v] do
        do sth
      end
    end
    size[u] <- size[u] + size[v]
  end
end func

如果认为do sth的复杂度是 \(O(T)\) 的话，那么上述伪代码的时间复杂度是 \(O(n^2T)\)，也就是说上述遍历整棵树外加 \(3\) 层 for 循环的时间复杂度是 \(O(n^2)\)。

证明上可以考虑如果记录每个点的 dfs 序，\(i\) 可以看作枚举 \(u\) 所有已经到达过的除 \(v\) 及其子树以外的后代。\(j\) 可以看作枚举 \(v\) 的后代。那么考虑任何一对点都只会在他们的 LCA \(u\) 处被枚举到，即每对点只会枚举一次，且 \(i\) 的 dfs 序一定小于 \(j\) 的 dfs 序，由于整个 dfs 序序列的顺序对个数又 \(O(n^2)\) 个，所以内层两层循环的总次数是 \(O(n^2)\) 的。于是上述代码的时间复杂度是 \(O(n^2T)\)，其中 do sth 的时间复杂度是 \(O(T)\)

回到这个题，最直接的DP是设 \(f_{u, i}\) 为以 \(u\) 为根的子树，选了 \(i\) 个黑点的最优答案，但是发现无法转移，因为全局最优解和局部最优没什么关系，于是考虑像 [NOI2019]回家的路 一样，对贡献去做DP。

考虑将上述状态更换为以 \(u\) 为根的子树，选了 \(i\) 个黑点对答案的最大贡献是多少，在转移时考虑每个子树的贡献即可。

具体的，设 \(f_{u, i, j}\) 为以 \(u\) 为根的子树，考虑前 \(i\) 个孩子，选了 \(j\) 个黑点的最大贡献。

那么转移显然：

\[f_{u,i,j} = \max_{h = 0}^{j} f_{u, i-1, h} + f_{v, size[v], j-h} + val
\]

其中 \(val\) 是这条边对答案造成的贡献，具体为边两侧黑点个数的乘积加上两侧白点个数的乘积的和再乘上边权。

例如子树中选择了 \(x\) 个黑点，那么贡献为

\[[x \times (k-x) + (size_v - x) \times (n - k - size_v + x)] \times e_w
\]

其中 \(e_w\) 为边权。

考虑上述转移中，\(u\) 的状态 \(i\) 只依赖于状态 \(i-1\) 和 \(size_v\)，因此可以滚动数组。滚动后的状态转移方程为：

\[f_{u,j} = \max_{h=0}^j f_{u, h} + f_{v, j-h} + val
\]

再转移时需要保证 \(f_{u,h}\) 是没有被 \(v\) 更新过的答案，也就是说比 \(j\) 小的状态应该在 \(j\) 之后更新，因此倒序枚举 \(j\) 即可。

在考虑最后一个问题：我们设计的方程是填表转移，即枚举了上面伪代码中 \(i\) 和 \(j\) 的和以及 \(i\)，只有这样才能保证方程中 \(j\) 的转移是单调的从而滚动数组，因此我们需要把上述伪代码改成外层枚举两数和，内层枚举除 \(v\) 以外的已遍历过的子树的形式。

func dfs(u):
  size[u] <- 1
  for v in clild[u] do
    dfs(v)
    size[u] <- size[u] + size[v]
    for sum = 1 : size[u] do
      for i = max(0, sum - size[v]) : size[v] do
        do sth
      end
    end
  end
end func

虽然写成这样的时间复杂度十分难以分析，但是枚举和再枚举其中一个的复杂度显然和枚举两个求和的复杂度相同，因此上述代码的时间复杂度也为 \(O(n^2T)\)

由于单次转移是 \(O(1)\) 的，因此算法的总时间复杂度 \(O(n^2)\)

Code

#include <cstdio>
#include <cstring>
#include <algorithm>

const int maxn = 2003;

int n, K, dK;
int sz[maxn];
ll frog[maxn][maxn];

struct Edge {
  int v, w;
  Edge *nxt;

  Edge(const int _v, const int _w, Edge *h) : v(_v), w(_w), nxt(h) {}
};
Edge *hd[maxn];

void dfs(const int u, const int fa);

int main() {
  freopen("1.in", "r", stdin);
  qr(n); qr(K); dK = n - K;
  for (int i = 1, u, v, w; i < n; ++i) {
    u = v = w = 0; qr(u); qr(v); qr(w);
    hd[u] = new Edge(v, w, hd[u]);
    hd[v] = new Edge(u, w, hd[v]);
  }
  dfs(1, 0);
  qw(frog[1][K], '\n', true);
  return 0;
}

void dfs(const int u, const int fa) {
  sz[u] = 1;
  memset(frog[u] + 2, -1, 16008);
  for (auto e = hd[u]; e; e = e->nxt) if (e->v != fa) {
    int v = e->v; dfs(v, u); sz[u] += sz[v];
    for (int i = std::min(sz[u], K); ~i; --i) {
      for (int  j = i, lim = std::max(0, i - sz[v]); j >= lim; --j) if (~frog[u][j]) {
        int k = i - j; if (frog[v][k] == -1) continue;
        ll val = (1ll * k * (K - k) + 1ll * (sz[v] - k) * (dK - sz[v] + k)) * e->w;
        frog[u][i] = std::max(frog[u][i], frog[u][j] + frog[v][k] + val);
      }
    }
  }
}

Summary

1、上面那种神奇的枚举方式遍历整棵树的时间复杂度是 \(O(n^2)\) 的

2、当问题不满足最优子结构时，可以考虑DP每个子问题对答案的贡献。

3、很多时候刷表法难以滚动数组，需要转化成填表法

appreciation

感谢 @DDOSvoid 大爷与我的讨论以及对我的启发