Mo's Algorithm on Tree (tree/mo_on_tree.hpp)

View this file on GitHub
View document part on GitHub
Last update: 2025-09-25 11:45:13+07:00
Include: #include "tree/mo_on_tree.hpp"

Description

Mo’s Algorithm on Tree processes offline path queries on trees efficiently by reordering queries to minimize the number of vertex additions/removals. Uses Euler Tour to convert tree paths to ranges.

Operations

MoOnTree<LG>()
- Constructs empty Mo on Tree instance
- Time complexity: $\mathcal O(1)$
template<class G> MoOnTree<LG>(G &g, int q, int root = 0)
- Constructs Mo on Tree instance for tree g with q queries and root root
- LG is the logarithm of the maximum tree depth
- Time complexity: $\mathcal O(n \log n)$
void add_query(int u, int v)
- Adds a path query from u to v
- Time complexity: $\mathcal O(1)$
template<class Eval, class Add, class Del> void run(Eval eval, Add add, Del del)
- Processes all queries
- eval(i) is called when query i should be evaluated
- add(v) / del(v) are called when adding/removing vertex v
- Time complexity: $\mathcal O(n \sqrt{q})$ where $q$ is the number of queries

Depends on

Mo's Algorithm (other/mo.hpp)

Code

#pragma once

#include "../other/mo.hpp"

template<int LG> struct MoOnTree {
    int n, qi = 0;
    vector<int> in, out, tour, w;
    vector<pair<int, int>> qs;
    array<vector<int>, LG> up;

    MoOnTree() {}
    template<class G>
    MoOnTree(G &g, int q, int root = 0) :
        n((int)g.size()), in(n), out(n), tour(2 * n), w(q), qs(q) {
        up.fill(vector<int>(n, -1));
        int t = 0;
        auto dfs = [&](auto &self, int v, int p = -1) -> void {
            up[0][v] = p;
            for (int l = 0; l + 1 < LG && up[l][v] != -1; l++)
                up[l + 1][v] = up[l][up[l][v]];
            tour[t] = v, in[v] = t++;
            for (int c : g[v]) {
                if (c == p) continue;
                self(self, c, v);
            }
            tour[t] = v, out[v] = t++;
        };
        dfs(dfs, root);
    }

    void add_query(int u, int v) {
        if (in[u] > in[v]) swap(u, v);
        w[qi] = lca(u, v);
        qs[qi] = {w[qi] == u ? in[u] : out[u], in[v] + 1};
        qi++;
    }

    template<class Eval, class Add, class Del>
    void run(Eval eval, Add add, Del del) {
        Mo mo(2 * n, qs);
        vector<bool> vis(n);
        auto mod = [&](int v) {
            v = tour[v], (vis[v] = !vis[v]) ? add(v) : del(v);
        };
        mo.run(
            [&](int i) {
                auto [u, v] = qs[i];
                u = tour[u], v = tour[v - 1];
                if (w[i] != u && w[i] != v) add(w[i]);
                eval(i);
                if (w[i] != u && w[i] != v) del(w[i]);
            },
            mod,
            mod
        );
    }

    bool anc(int p, int v) const {
        return p == -1 || (in[p] <= in[v] && out[v] <= out[p]);
    }
    int lca(int u, int v) {
        if (anc(u, v)) return u;
        if (anc(v, u)) return v;
        for (int l = LG - 1; l >= 0; l--)
            if (!anc(up[l][u], v)) u = up[l][u];
        return up[0][u];
    };
};

#line 2 "tree/mo_on_tree.hpp"

#line 2 "other/mo.hpp"

struct Mo {
    int q;
    const int B;
    vector<int> qi;
    vector<pair<int, int>> qs;

    int l = 0, r = 0;

    Mo(int n, const vector<pair<int, int>> &_qs) :
        q((int)_qs.size()), B(blk_sz(n, q)), qi(q), qs(_qs) {
        iota(begin(qi), end(qi), 0);
        sort(begin(qi), end(qi), [&](int i, int j) {
            auto [li, ri] = qs[i];
            auto [lj, rj] = qs[j];
            int bi = li / B, bj = lj / B;
            if (bi != bj) return bi < bj;
            if (ri != rj) return (bi & 1) != (ri < rj);
            return li < lj;
        });
    }

    template<class Eval, class AL, class DL, class AR, class DR>
    void run(Eval eval, AL add_l, DL del_l, AR add_r, DR del_r) {
        for (int i : qi) {
            auto [ql, qr] = qs[i];
            while (ql < l) add_l(--l);
            while (r < qr) add_r(r++);
            while (l < ql) del_l(l++);
            while (qr < r) del_r(--r);
            eval(i);
        }
    }
    template<class Eval, class Add, class Del>
    void run(Eval eval, Add add, Del del) {
        run(eval, add, del, add, del);
    }

  private:
    static int blk_sz(int n, int q) {
        return max(1, int(n / max<double>(1, sqrt(q * 2.0 / 3.0))));
    }
};
#line 4 "tree/mo_on_tree.hpp"

template<int LG> struct MoOnTree {
    int n, qi = 0;
    vector<int> in, out, tour, w;
    vector<pair<int, int>> qs;
    array<vector<int>, LG> up;

    MoOnTree() {}
    template<class G>
    MoOnTree(G &g, int q, int root = 0) :
        n((int)g.size()), in(n), out(n), tour(2 * n), w(q), qs(q) {
        up.fill(vector<int>(n, -1));
        int t = 0;
        auto dfs = [&](auto &self, int v, int p = -1) -> void {
            up[0][v] = p;
            for (int l = 0; l + 1 < LG && up[l][v] != -1; l++)
                up[l + 1][v] = up[l][up[l][v]];
            tour[t] = v, in[v] = t++;
            for (int c : g[v]) {
                if (c == p) continue;
                self(self, c, v);
            }
            tour[t] = v, out[v] = t++;
        };
        dfs(dfs, root);
    }

    void add_query(int u, int v) {
        if (in[u] > in[v]) swap(u, v);
        w[qi] = lca(u, v);
        qs[qi] = {w[qi] == u ? in[u] : out[u], in[v] + 1};
        qi++;
    }

    template<class Eval, class Add, class Del>
    void run(Eval eval, Add add, Del del) {
        Mo mo(2 * n, qs);
        vector<bool> vis(n);
        auto mod = [&](int v) {
            v = tour[v], (vis[v] = !vis[v]) ? add(v) : del(v);
        };
        mo.run(
            [&](int i) {
                auto [u, v] = qs[i];
                u = tour[u], v = tour[v - 1];
                if (w[i] != u && w[i] != v) add(w[i]);
                eval(i);
                if (w[i] != u && w[i] != v) del(w[i]);
            },
            mod,
            mod
        );
    }

    bool anc(int p, int v) const {
        return p == -1 || (in[p] <= in[v] && out[v] <= out[p]);
    }
    int lca(int u, int v) {
        if (anc(u, v)) return u;
        if (anc(v, u)) return v;
        for (int l = LG - 1; l >= 0; l--)
            if (!anc(up[l][u], v)) u = up[l][u];
        return up[0][u];
    };
};