[学习笔记]换根 DP 的常用处理方式

换根 DP，又称作二次扫描法，通常用于“对每个根求出树形 DP 的答案”。以每个点作为根节点进行一次树形 DP 的代价通常无法承受，因此我们会使用两次 DFS：

• 第一次 DFS 指定一个点为根节点，运行一次常规的树形 DP。
• 第二遍 DFS 进行换根 DP，得到将根转移到每个节点的答案。

换根 DP 常用的例题是一些具有特别良好性质的问题，例如求所有节点的深度总和。当根节点从 \(u\) 转移到 \(v\) 时，所有在 \(v\) 子树中的节点（以 \(0\) 为根/以 \(u\) 为根）深度减一，而其余节点深度加一。在第一遍 DFS 中处理出子树大小，在第二遍 DFS 中即可直接用 \(ans_v = ans_u + n - sz_v - sz_v = ans_u + n - 2 sz_v\) 更新答案，甚至没有用到第一遍树形 DP 的“子树节点深度和”信息。

但对于一般的问题，我们并不总能如此简单地更新答案。例如求最深节点的深度，不另外使用数据结构，无法快速地完成子树深度加减、全局深度最大值这两个操作。

算法步骤

一般地，在第二次 DFS 搜索到 \(u\) 时，首先更新以 \(u\) 为根的答案；之后枚举 \(u\) 的子节点 \(v\)，要将根从 \(u\) 转移到 \(v\)，我们分为以下几步：

1. 将 \(v\) 的贡献从 \(dp_u\) 移除；
2. 把 \(u\) 作为 \(v\) 的子节点，将贡献加入到 \(dp_v\)。
3. 以 \(v\) 为新的根节点继续进行第二次扫描。
4. 回滚 \(dp_u\) 的状态。

注意第三步结束后，我们需要把 \(dp_u\) 的状态恢复到第一步之前，因为还要对其它的儿子重复这一个过程。而在 \(v\) 的（以 \(0\) 为根/以 \(u\) 为根）子树上的节点均已完成所有计算，它们的状态并不重要。说到回滚，就自然有多种方式。最简单地，若 Copy 的代价可以承担（例如要维护状态所占空间为 \(O(1)\) 时），我们可以每次为 \(dp_u\) 创建一个新的副本，在其上执行前两步操作。

让我们回到所有节点深度总和的问题，有 \(dp_u = 1 +\sum_{v \in sons(u)} (sz_v + dp_v), sz_u = 1 +\sum_{v \in sons(u)} sz_v\)。换根时我们依次执行上面的步骤，并使用“创建副本”的方法而避免回滚。

1. 创建新的副本，在其中去掉 \(v\) 的贡献：
   • \(dp_u' = dp_u - sz_v - dp_v, sz_u' = sz_u - sz_v = n - sz_v\)。
2. 将 \(u\) 作为 \(v\) 的子节点，将贡献加入到 \(dp_v\)：
   • \(dp_v \larr dp_v + sz_u' + dp_u' = dp_u + n - 2sz_v, sz_v \larr sz_v + sz_u' = n\)。

虽然推导步骤不同，我们得到了相同的式子。

但复制的代价并不总是能够承担。例如有的时候我们被迫（或为了方便）需要使用一个数据结构（如 STL set/ multiset）维护所有子树的信息以进行 DP 转移，此时就需要采用另外的回滚方法。例如：

• 可以记录所有操作，并逐个进行逆操作。它可以是各种层面上的逆操作。
  • 如果所有宏观上的操作都是可逆的，可以倒序执行一遍所有操作的逆操作。例如 multiset 的 insert 与 extract；以及更高层面的，题目中具体使用的操作。
    • 甚至更暴力地，对每一次内存写入，记录位置以及写入前的值；回滚时逆序撤销。
  • 尽管 \(dp_v\) 已经不再需要用到，但出于正确回滚 \(dp_u\) 的需要，\(dp_v\) 很可能会被一并回滚。
• 甚至可以采用可持久化等方式。

例如 CF1796E 题，为了方便使用 multiset 维护最小值、次小值。建立宏观上互为逆操作的的添加/删除贡献的函数 add 与 del，以恰当的顺序和参数调用，以确保贡献的加入/移除、状态的回滚都能够正确完成。参考 这份题解 的相关代码：

void add(int u, int val) {
    if (f[u].size() >= 2) se.erase(se.find(*next(f[u].begin())));
    f[u].insert(val);
    if (f[u].size() >= 2) se.insert(*next(f[u].begin()));
}

void del(int u, int val) {
    if (f[u].size() >= 2) se.erase(se.find(*next(f[u].begin())));
    f[u].erase(f[u].find(val));
    if (f[u].size() >= 2) se.insert(*next(f[u].begin()));
}

void dfs2(int u, int fa) {
    ans = max(ans, min(getlen(u), se.empty() ? INF : *se.begin()));
    for (auto v : G[u]) {
        if (v == fa) continue;
        del(u, getlen(v));
        add(v, getlen(u));
        dfs2(v, u);
        del(v, getlen(u));
        add(u, getlen(v));
    }
}

有关贡献的移除

尽管我们可能需要回滚第二步“把 \(u\) 的贡献添加到 \(dp_v\)”，它在效果上等价于移除 \(u\) 对 \(dp_v\) 的贡献，一般情况下贡献的移除不像回滚这么简单。这是因为我们等价于要计算如下问题：

• 对 \(u\) 的每个儿子 \(v\)，计算 \(fa_u\)（若存在）以及 \(u\) 除 \(v\) 外的所有儿子的贡献的聚合。

对加减乘除等有结合律且存在逆元的聚合操作，移除贡献特别容易。但一般情形下，移除一个特定节点的贡献比回滚最后一次操作要难得多。

回到最深节点深度问题，有 \(dp_u = \max\limits_{v \in sons(u)} \{1 +dp_v\}\)。第一眼看上去 \(\max\) 操作是不可逆的，因为它始终只保存当前最大的贡献，会丢失信息。但注意到所谓的“移除”永远只会涉及到一个节点，这是个很重要的性质：我们只要在更新 \(dp_u\) 的同时记录次大值 \(sdp_u\)，按照 \(dp_v\) 是否是最大值选择 \(dp_u'\)。按照上面的步骤，算法如下：

• 若 \(dp_u = dp_v + 1\)，\(dp_u' = sdp_u\)；否则 \(dp_u' = dp_u\)。
• 根据 \(dp_u' + 1\) 更新 \(dp_v\) 与 \(sdp_v\)。

同样我们使用复制副本以避免回滚。类似的，我们可以 \(O(1)\) 处理最小值、次小值等贡献的移除。

对更特殊的聚合操作，如 \(\gcd\)、\({\rm lcm}\)、对非质数的模乘法，这个性质也帮不上大忙。所幸我们有熟悉的小技巧：前后缀分解（我们总认为聚合操作是满足交换律的。若不满足，也不代表一定不能使用这一技巧），在状态的大小、聚合的均为 \(O(1)\) 时，总能保证换根 DP 的总时间复杂度为 \(O(n)\)。