Black White Tree

题意¶

有一棵 $n$ 个点的树以及一个参数 $k$，结点依次标号 $1, \dots, n$。结点 $i$ 有权值 $w_i$，其深度为 $d_i$。

初始时每个结点均为白色，可以选某些点染黑，但需要满足以下条件：

称黑色结点 $u, v$ 是连通的当且仅当 $u$ 到 $v$ 路径上的结点均为黑色；
对于任意连通结点 $u, v$，均要求 $|d_u-d_v|\leq k$。

最大化黑色结点权值之和。

$1\leq n \leq 10^6$，$0\leq k\leq 10^6$，$0\leq w_i\leq 10^9$。

题目来源¶

@Serval

题解¶

Hint 1

$k = 0$ 的特例是最大独立集。

Hint 2

$k = 1$ 的特例也可以沿用 $k = 0$ 的做法 DP 解决。

Hint 3

设 $f_{u,i}$ 表示当前考虑以 $u$ 为根的子树，$u$ 联通的结点与 $u$ 深度差最大为 $i - 1$ 时黑色结点权值和的最大值。特殊地，$f_{u,0}$ 表示结点 $u$ 为白色。直接 DP 是 $O(nk)$ 的，有什么办法优化吗？

Hint 4

一种优化树上 DP 的方式是长链剖分。

Hint 5

设 $v$ 取遍 $u$ 的子结点，容易写出状态转移方程： $$ \begin{aligned} f_{u,0} &= \sum_v \max_{0\leq j\leq k+1} {f_{v,j}} \\ f_{u,i} &= w_u + \sum_v \max_{0\leq j\leq i-1} {f_{v,j}} \qquad (i > 0) \end{aligned} $$ 转移时需要取前缀最大值，如何优化？

Hint 6

不妨设 $g_{u,i} = \max_{0\leq j\leq i} \{f_{u,j}\}$，有什么效果吗？

Hint 7

我们只需要维护 $g$ 而不需要维护 $f$ 即可完成 DP！ $$ \begin{aligned} g_{u,0} &= \sum_v g_{v,k+1} {f_{v,j}} \\ g_{u,i} &= w_u + \sum_v g_{v,i-1} \qquad (i > 0) \end{aligned} $$ 但是考虑这样一种情况：长度为 $k$ 与长度为 $1$ 的两条链合并，长度为 $1$ 的链对 $g$ 均有贡献。

Hint 8

根据 $g$ 的定义，$g_{u,i}$ 随 $i$ 递增单调不减。

Hint 9

可以对单调序列作差分，通过维护差分序列而维护单调序列。

Solution

定义 $g_{u,i}=\max_{0\leq j\leq i} f_{u,i}$，$g_{u,i}$ 单调不降是显然的。考虑转移时需要的操作：

查询单调序列最大值（查询 $g_{u,k+1}$）
删除单调序列末尾元素（删去 $g_{u,k+2}$ 以保证单调序列中所有值合法）
合并两个单调序列，对应位置相加（轻儿子 DP 值向父结点合并）
单调序列整体加一个非负数（对于 $0<i\leq k+1$ 的 $g_{u,i}$ 加上 $w_u$）
单调序列前端添加元素，并对后面所有项取 $\max$（加入 $g_{u,0}$）

可以对单调序列作差分，并维护差分序列非 $0$ 的项。额外维护差分序列和即可查询序列最大值，删除差分序列的末尾元素即可删除单调序列末尾元素，差分序列首项加法即可单调序列整体加法，这一部分时间复杂度 $O(1)$。

合并两个单调序列只需提取两个差分序列的公共前缀，双指针合并后再与非公共后缀拼接。两个长度为 $l_1,l_2$ 的单调序列合并的时间复杂度 $O(\min\{l_1,l_2\})$。

在单调序列前端添加元素需要在差分序列前端添加元素，并删除或修改后面连续一段的差分值以取 $\max$，由于一个结点最多加入一个差分值（$g_{u,0}$），这一部分均摊时间复杂度 $O(1)$。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。