每日算法 — 使用java实现树上启发式合并：子树统计的DSU on Tree优化

算法概述

树上启发式合并（DSU on Tree，又称Sack on Tree）是一种高效的树形动态规划优化技术。它巧妙地利用”小合并到大”的思想，在 $O(n \log n)$ 时间内解决大量子树统计问题，如子树颜色种类数、子树中众数出现次数等。相比传统的暴力 $O(n^2)$ 解法，DSU on Tree将复杂度降低了一个数量级。

核心思想

DSU on Tree 的核心在于轻重链剖分与保留重儿子信息的结合：

轻重链剖分：为每个节点计算子树大小，子树最大的儿子称为”重儿子”，其余为”轻儿子”
先处理轻儿子，再处理重儿子：轻儿子的贡献计算后清空，重儿子的贡献保留
小合并到大：将轻儿子的统计信息暴力合并到重儿子的统计容器中

通过这种策略，每个节点的信息最多被向上合并 $O(\log n)$ 次，总复杂度为 $O(n \log n)$。

算法步骤

第一遍DFS：计算每个节点的子树大小 sz[u]，确定重儿子 heavy[u]
第二遍DFS：
先递归处理所有轻儿子（不保留统计信息）
再递归处理重儿子（保留统计信息）
将当前节点和所有轻儿子的贡献加入统计容器
回答以当前节点为根的子树查询
如果是轻儿子子树，清空统计信息

Java 完整实现

import java.io.*;
import java.util.*;

/**
 * 树上启发式合并（DSU on Tree）完整实现
 * 解决子树颜色种类统计问题
 */
public class DSUOnTree {

    // 邻接表存树
    static List<Integer>[] tree;
    // 每个节点的颜色
    static int[] color;
    // 子树大小
    static int[] sz;
    // 重儿子
    static int[] heavy;
    // 答案数组，ans[u]表示以u为根的子树中颜色种类数
    static int[] ans;

    // 统计容器：cnt[c]表示颜色c在当前统计范围内的出现次数
    static int[] cnt;
    // 当前统计范围内不同颜色的数量
    static int currentDistinct;

    // 是否需要保留当前子树的统计信息（重儿子保留）
    static boolean keep;

    /**
     * 第一遍DFS：计算子树大小和重儿子
     */
    static void dfs1(int u, int parent) {
        sz[u] = 1;
        int maxSubSize = 0;

        for (int v : tree[u]) {
            if (v != parent) {
                dfs1(v, u);
                sz[u] += sz[v];
                if (sz[v] > maxSubSize) {
                    maxSubSize = sz[v];
                    heavy[u] = v;
                }
            }
        }
    }

    /**
     * 添加或删除一个子树的所有节点贡献
     * @param u 当前节点
     * @param parent 父节点
     * @param delta +1表示添加，-1表示删除
     */
    static void add(int u, int parent, int delta) {
        // 处理当前节点
        int c = color[u];
        if (delta == 1) {
            if (cnt[c] == 0) currentDistinct++;
            cnt[c]++;
        } else {
            cnt[c]--;
            if (cnt[c] == 0) currentDistinct--;
        }

        // 递归处理所有子节点
        for (int v : tree[u]) {
            if (v != parent && v != heavy[u]) {
                add(v, u, delta);
            }
        }
    }

    /**
     * 第二遍DFS：DSU on Tree核心逻辑
     * @param u 当前节点
     * @param parent 父节点
     * @param keepFlag 是否保留当前子树的统计信息
     */
    static void dfs2(int u, int parent, boolean keepFlag) {
        // 1. 先处理所有轻儿子（不保留）
        for (int v : tree[u]) {
            if (v != parent && v != heavy[u]) {
                dfs2(v, u, false);
            }
        }

        // 2. 处理重儿子（保留）
        if (heavy[u] != -1) {
            dfs2(heavy[u], u, true);
        }

        // 3. 将当前节点和所有轻儿子的贡献加入统计
        int c = color[u];
        if (cnt[c] == 0) currentDistinct++;
        cnt[c]++;

        for (int v : tree[u]) {
            if (v != parent && v != heavy[u]) {
                add(v, u, 1); // 添加轻儿子子树的所有节点
            }
        }

        // 4. 回答当前子树的查询
        ans[u] = currentDistinct;

        // 5. 如果不保留，清空当前子树的所有统计信息
        if (!keepFlag) {
            add(u, parent, -1);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        // 读入节点数
        int n = Integer.parseInt(br.readLine().trim());

        // 初始化
        tree = new ArrayList[n + 1];
        for (int i = 1; i <= n; i++) tree[i] = new ArrayList<>();
        color = new int[n + 1];
        sz = new int[n + 1];
        heavy = new int[n + 1];
        ans = new int[n + 1];
        Arrays.fill(heavy, -1);

        // 读入颜色（假设颜色范围在1到n之间）
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            color[i] = Integer.parseInt(st.nextToken());
        }

        // 读入边
        for (int i = 0; i < n - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            tree[u].add(v);
            tree[v].add(u);
        }

        // 初始化统计容器（假设颜色最大值为n）
        cnt = new int[n + 1];
        currentDistinct = 0;

        // 执行DSU on Tree
        dfs1(1, -1);
        dfs2(1, -1, false);

        // 输出每个节点的子树颜色种类数
        StringBuilder sb = new StringBuilder();
        for (int i = 1; i <= n; i++) {
            sb.append(ans[i]).append(" ");
        }
        System.out.println(sb.toString().trim());
    }
}

复杂度分析

时间复杂度：$O(n \log n)$。每个节点的颜色信息最多被向上合并 $O(\log n)$ 次（因为每次合并到更大的子树，子树大小至少翻倍）
空间复杂度：$O(n)$。统计容器大小与颜色范围相关，递归栈深度为树高

适用场景

DSU on Tree 适用于以下类型的子树查询问题：
– 子树中不同颜色/值的种类数
– 子树中出现次数最多的颜色及其出现次数
– 子树中满足某种条件的节点数量
– 任何可以通过暴力添加/删除节点来维护统计信息的问题

扩展应用

通过修改 add 函数中的统计逻辑，DSU on Tree 可以灵活适配多种子树统计需求：
– 子树众数：维护当前出现次数的最大值
– 子树第k大：结合值域分块或线段树
– 子树异或和：维护异或值的计数

DSU on Tree 是处理静态子树查询问题的一把利器，在竞赛和工程实践中都有广泛应用。

Tagged with: DSU on Tree Java 子树统计树上启发式合并树形DP 轻重链剖分