每日算法 — 使用java实现树上莫队：离线查询与欧拉序分块优化

可莉爸爸 AI算法

一、问题背景

在树上路径查询问题中，我们经常需要回答形如”节点 $u$ 到 $v$ 的路径上有多少个黑点？”这样的离线查询。直接每条查询都用BFS/DFS遍历路径会很慢。

树上莫队 是莫队算法在树结构上的优雅扩展。它通过 欧拉序 将树上路径问题转化为线性区间的查询，复用了莫队的分块优化思想。

二、核心思想

2.1 欧拉序

对树进行DFS，记录每个节点在DFS过程中进入和离开的次序：

  • 第一次访问节点 $u$ 时记录 in[u]
  • 离开时记录 out[u]

这样在欧拉序中，从根到 $u$ 的路径就是 in[root] 到 in[u] 的一个区间。

2.2 路径到区间的转化

对于路径 $u \to v$：
– 找到 $u, v$ 的最近公共祖先 $lca$
– 路径上的所有节点对应欧拉序中的区间 $[in[lca]+1, in[v]]$ 中的部分（若 $lca$ 不在端点）

巧妙地处理：路径上 $lca$ 出现 偶数次（一次入一次出），其它节点出现 偶数次（最终只算一次）。通过维护节点访问次数的奇偶性，配合 LCA 求解器，我们可以高效回答路径查询。

三、Java实现

3.1 树结构与欧拉序生成

import java.util.*;

/**
 * 树上莫队算法实现
 * 用于回答形如"路径(u,v)上的某种统计"问题
 */
public class TreeMoAlgorithm {

    private final int n;           // 节点数
    private final List<Integer>[] tree;  // 邻接表
    private final int[] in;        // 欧拉序位置
    private final int[] out;
    private final int[] euler;     // 欧拉序数组
    private final int[] nodeAt;    // 欧拉序中每个位置对应的节点
    private final int[] st;        // RMQ稀疏表（用于LCA）
    private final int[] depth;
    private int time = 0;

    @SuppressWarnings("unchecked")
    public TreeMoAlgorithm(int n) {
        this.n = n;
        this.tree = new ArrayList[n];
        for (int i = 0; i < n; i++) tree[i] = new ArrayList<>();
        in = new int[n];
        out = new int[n];
        euler = new int[2 * n];
        nodeAt = new int[2 * n];
        depth = new int[n];
    }

    public void addEdge(int u, int v) {
        tree[u].add(v);
        tree[v].add(u);
    }

    /**
     * DFS生成欧拉序
     */
    public void generateEuler(int root) {
        time = 0;
        dfs(root, -1);
        buildLcaTable();
    }

    private void dfs(int u, int parent) {
        in[u] = time;
        euler[time] = u;
        nodeAt[time] = u;
        time++;

        for (int v : tree[u]) {
            if (v == parent) continue;
            depth[v] = depth[u] + 1;
            dfs(v, u);
        }

        out[u] = time;
        euler[time] = u;
        nodeAt[time] = u;
        time++;
    }

3.2 LCA稀疏表

    private int[][] stTable;

    private void buildLcaTable() {
        int len = euler.length;
        int k = (int) Math.ceil(Math.log(len) / Math.log(2)) + 1;
        stTable = new int[k][len];
        stTable[0] = euler.clone();
        for (int i = 1; (1 << i) <= len; i++) {
            for (int j = 0; j + (1 << i) <= len; j++) {
                if (depth[stTable[i-1][j]] < depth[stTable[i-1][j + (1<<(i-1))]]) {
                    stTable[i][j] = stTable[i-1][j];
                } else {
                    stTable[i][j] = stTable[i-1][j + (1<<(i-1))];
                }
            }
        }
    }

    /**
     * 使用RMQ求最近公共祖先
     */
    public int lca(int u, int v) {
        if (in[u] > in[v]) {
            int tmp = u; u = v; v = tmp;
        }
        int l = in[u], r = in[v];
        if (l > r) {
            int tmp = l; l = r; r = tmp;
        }
        int k = (int) Math.floor(Math.log(r - l + 1) / Math.log(2));
        if (depth[stTable[k][l]] < depth[stTable[k][r - (1<<k) + 1]]) {
            return stTable[k][l];
        } else {
            return stTable[k][r - (1<<k) + 1];
        }
    }

3.3 莫队查询结构

    /**
     * 莫队查询
     * 包含 l, r, lca, 查询id
     */
    static class Query {
        int l, r, lca, id;
        int block;

        Query(int l, int r, int lca, int id) {
            this.l = l;
            this.r = r;
            this.lca = lca;
            this.id = id;
        }
    }

    /**
     * 树上的路径查询
     * 转化为欧拉序上的区间查询
     */
    public int[] processQueries(int[][] queries) {
        int q = queries.length;
        Query[] mq = new Query[q];

        // 块大小：通常取 O(sqrt(n))
        int blockSize = (int) Math.sqrt(time) + 1;

        for (int i = 0; i < q; i++) {
            int u = queries[i][0], v = queries[i][1];
            int w = lca(u, v);

            if (in[u] > in[v]) {
                int tmp = u; u = v; v = tmp;
            }

            int l, r;
            if (u == w) {
                // lca 是其中一个端点
                l = in[u];
                r = in[v];
                mq[i] = new Query(l, r, -1, i);
            } else {
                // lca 不在端点，需要额外处理lca
                l = out[u];
                r = in[v];
                mq[i] = new Query(l, r, w, i);
            }
            mq[i].block = l / blockSize;
        }

        // 莫队排序
        Arrays.sort(mq, (a, b) -> {
            if (a.block != b.block) return Integer.compare(a.block, b.block);
            if ((a.block & 1) == 1) return Integer.compare(a.r, b.r);
            return Integer.compare(b.r, a.r);
        });

        // 离线处理
        int[] results = new int[q];
        int[] color = new int[n];  // 节点当前是否被激活
        int[] count = new int[n];  // 颜色计数
        int totalColors = 0;        // 当前不同颜色数
        int curL = 0, curR = -1;

        for (Query query : mq) {
            // 扩展或收缩窗口
            while (curL > query.l) addNode(euler[--curL], color, count);
            while (curR < query.r) addNode(euler[++curR], color, count);
            while (curL < query.l) removeNode(euler[curL++], color, count);
            while (curR > query.r) removeNode(euler[curR--], color, count);

            // 单独处理lca
            int ans = totalColors;
            if (query.lca != -1) {
                if (color[query.lca] == 0) ans++;
            }
            results[query.id] = ans;
        }

        return results;
    }

    /**
     * 激活/取消激活一个节点
     */
    private void addNode(int node, int[] color, int[] count) {
        color[node] ^= 1;
        int c = node;  // 这里可以扩展为按颜色分组
        if (color[node] == 1) {
            count[c]++;
            if (count[c] == 1) totalColors++;
        } else {
            count[c]--;
            if (count[c] == 0) totalColors--;
        }
    }

    private void removeNode(int node, int[] color, int[] count) {
        addNode(node, color, count);  // 异或操作对称
    }

3.4 完整测试

    public static void main(String[] args) {
        System.out.println("=== 树上莫队算法演示 ===\n");

        // 构建示例树
        //        0
        //       /|\
        //      1 2 3
        //     /|   |
        //    4 5   6
        //   /|
        //  7 8

        TreeMoAlgorithm solver = new TreeMoAlgorithm(9);
        solver.addEdge(0, 1);
        solver.addEdge(0, 2);
        solver.addEdge(0, 3);
        solver.addEdge(1, 4);
        solver.addEdge(1, 5);
        solver.addEdge(3, 6);
        solver.addEdge(4, 7);
        solver.addEdge(4, 8);

        solver.generateEuler(0);

        // 测试LCA
        System.out.println("=== LCA 测试 ===");
        System.out.println("LCA(7, 8) = " + solver.lca(7, 8));
        System.out.println("LCA(5, 6) = " + solver.lca(5, 6));
        System.out.println("LCA(7, 3) = " + solver.lca(7, 3));

        // 路径查询
        System.out.println("\n=== 路径查询测试 ===");
        int[][] queries = {
            {7, 5},   // 7-1-5
            {7, 8},   // 7-4-8
            {5, 6},   // 5-1-0-3-6
            {4, 2},   // 4-1-0-2
            {0, 0}    // 单点
        };

        int[] results = solver.processQueries(queries);
        for (int i = 0; i < queries.length; i++) {
            System.out.println("路径查询 " + 
                             queries[i][0] + " -> " + queries[i][1] + 
                             " 路径长度：" + results[i]);
        }

        // 性能对比（可选）
        System.out.println("\n=== 性能测试 ===");
        Random rand = new Random(42);
        int n = 1000;
        TreeMoAlgorithm bigSolver = new TreeMoAlgorithm(n);
        for (int i = 1; i < n; i++) {
            bigSolver.addEdge(i, rand.nextInt(i));
        }
        bigSolver.generateEuler(0);

        int[][] bigQueries = new int[1000][2];
        for (int i = 0; i < 1000; i++) {
            int u = rand.nextInt(n);
            int v = rand.nextInt(n);
            bigQueries[i][0] = u;
            bigQueries[i][1] = v;
        }

        long start = System.currentTimeMillis();
        bigSolver.processQueries(bigQueries);
        long elapsed = System.currentTimeMillis() - start;
        System.out.println("1000个查询耗时：" + elapsed + " ms");
    }
}

四、复杂度分析

操作 时间复杂度 空间复杂度
欧拉序生成 $O(n)$ $O(n)$
LCA稀疏表 $O(n \log n)$ $O(n \log n)$
排序查询 $O(q \log q)$ $O(q)$
莫队处理 $O((n + q) \sqrt{n})$ $O(n)$

其中 $n$ 为节点数，$q$ 为查询数。

相比朴素的 $O(q \cdot n)$ 实现，树上莫队在处理大量离线查询时性能提升巨大。

五、变体与扩展

  • 回滚莫队：处理难以撤销的操作时使用栈+回滚
  • 二次离线莫队：进一步降低复杂度到 $O(n\sqrt{q})$
  • 带修改的树上莫队：结合时间维度，处理动态更新
  • 树上分块：对查询较少的场景，树上分块是替代方案

树上莫队是处理大规模树路径离线查询问题的标准武器，灵活掌握它能够应对很多高级算法竞赛题目。

Tagged with: