m1une's library

This documentation is automatically generated by online-judge-tools/verification-helper

View on GitHub

:heavy_check_mark: Mo on Tree
(graph/tree/mo_on_tree.hpp)

Overview

MoOnTree<T> reorders static path queries on a tree and maintains the current path through user-defined insertion and deletion callbacks. It is useful for path statistics such as the number of distinct colors, value frequencies, and other set-like aggregates whose one-vertex update is fast.

The structure uses a doubled Euler tour. Every vertex occurs once on entry and once on exit, so the parity of an Euler interval describes one tree path. A lowest common ancestor outside that interval is inserted temporarily while its query is answered.

The graph must be a connected undirected tree built with Graph<T>::add_edge. Inactive edges are ignored by the underlying tree traversal. All indices are zero-based.

Construction

MoOnTree();
MoOnTree(const m1une::graph::Graph<T>& graph, int root = 0);
void build(const m1une::graph::Graph<T>& graph, int root = 0);

Construction roots the tree, builds heavy-light LCA metadata, and creates a doubled Euler tour in $O(N)$ time and memory. build discards all registered queries.

The chosen root does not change which vertices or edges belong to a path. It only determines the child vertex used to represent each edge.

Registering Queries

int add_query(int from, int to);
int add_edge_query(int from, int to);

add_query registers the inclusive vertex path from from to to. add_edge_query registers the edges on that path. Both return an insertion-order query ID and take $O(\log N)$ time for an LCA query.

For an edge query, callbacks receive the endpoint farther from the chosen root. Thus the corresponding original edge ID is parent_edge(vertex). The empty path registered by add_edge_query(v, v) contains no edge.

Vertex queries and edge queries cannot be mixed in one collection. Use two MoOnTree objects if both kinds are needed. clear() removes all queries and allows either kind to be selected again.

Running the Algorithm

mo.run(add, remove, answer);
mo.run(add, remove, answer, block_size);

The callbacks have these meanings:

In edge mode, every callback vertex has a valid nonnegative parent_edge(vertex); transitions involving the dummy root item are suppressed.

Callbacks can be invoked many times for one item. After run returns, the maintained state is not guaranteed to be empty. For a query whose LCA needed a temporary insertion, that insertion is removed immediately after answer.

With callback cost $O(F)$, the automatic block size gives the usual $O((N\sqrt Q+Q)F+Q\log Q)$ time bound. A positive explicit block size can be used for problem-specific tuning.

Methods

Method Description Complexity
MoOnTree() Creates an empty object. $O(1)$
MoOnTree(graph, root) Builds the rooted tree and doubled tour. $O(N)$
build(graph, root) Rebuilds and clears all queries. $O(N)$
size() Returns the number of vertices. $O(1)$
empty() Returns whether the tree is empty. $O(1)$
query_count() Returns the number of registered queries. $O(1)$
queries() Returns registered queries in insertion order. $O(1)$
parent(v) Returns the rooted parent, or -1 at the root. $O(1)$
parent_edge(v) Returns the edge to the parent, or -1 at the root. $O(1)$
depth(v) Returns the unweighted rooted depth. $O(1)$
lca(u, v) Returns the lowest common ancestor. $O(\log N)$
reserve(q) Reserves space for q queries. $O(Q)$
clear() Removes registered queries. $O(Q)$
add_query(u, v) Adds the inclusive vertex path. Amortized $O(\log N)$
add_edge_query(u, v) Adds the edge path. Amortized $O(\log N)$
order(block_size) Returns query IDs in Mo order. $O(Q\log Q)$
run(...) Processes all queries. See above

Public root, entry, exit, and tour expose the chosen root and doubled Euler tour. For every vertex v, tour[entry[v]] and tour[exit[v]] both equal v.

Each public Query stores from, to, transformed half-open interval [left, right), optional extra LCA, insertion-order id, and the edge mode flag.

Example

This computes the number of distinct colors on every vertex path:

#include "graph/graph.hpp"
#include "graph/tree/mo_on_tree.hpp"

#include <vector>

int main() {
    m1une::graph::Graph<int> graph(5);
    graph.add_edge(0, 1);
    graph.add_edge(0, 2);
    graph.add_edge(1, 3);
    graph.add_edge(1, 4);

    std::vector<int> color = {0, 1, 0, 2, 1};
    m1une::tree::MoOnTree<int> mo(graph);
    mo.add_query(3, 2);
    mo.add_query(3, 4);

    std::vector<int> frequency(3);
    std::vector<int> result(mo.query_count());
    int distinct = 0;
    mo.run(
        [&](int vertex) {
            if (frequency[color[vertex]]++ == 0) distinct++;
        },
        [&](int vertex) {
            if (--frequency[color[vertex]] == 0) distinct--;
        },
        [&](int query_id) {
            result[query_id] = distinct;
        }
    );
}

Depends on

Required by

Verified with

Code

#ifndef M1UNE_TREE_MO_ON_TREE_HPP
#define M1UNE_TREE_MO_ON_TREE_HPP 1

#include <algorithm>
#include <cassert>
#include <vector>

#include "../../algo/offline/mo.hpp"
#include "../graph.hpp"
#include "heavy_light_decomposition.hpp"

namespace m1une {
namespace tree {

// Offline Mo's algorithm for static paths in a tree.
template <class T = int>
struct MoOnTree {
    struct Query {
        int from;
        int to;
        int left;
        int right;
        int extra;
        int id;
        bool edge;
    };

    int root;
    std::vector<int> entry;
    std::vector<int> exit;
    std::vector<int> tour;

   private:
    int _n;
    HeavyLightDecomposition<T> _hld;
    m1une::algo::Mo _mo;
    std::vector<Query> _queries;

    void check_vertex(int vertex) const {
        assert(0 <= vertex && vertex < _n);
        assert(entry[vertex] != -1);
    }

    int add_path_query(int from, int to, bool edge) {
        check_vertex(from);
        check_vertex(to);
        assert(_queries.empty() || _queries.front().edge == edge);
        int original_from = from;
        int original_to = to;
        if (entry[from] > entry[to]) std::swap(from, to);

        int ancestor = _hld.lca(from, to);
        int left;
        int right = entry[to] + 1;
        int extra = -1;
        if (ancestor == from) {
            left = entry[from] + int(edge);
        } else {
            left = exit[from];
            if (!edge) extra = ancestor;
        }

        int id = _mo.add_query(left, right);
        _queries.push_back(Query{
            original_from,
            original_to,
            left,
            right,
            extra,
            id,
            edge,
        });
        return id;
    }

   public:
    MoOnTree() : root(-1), _n(0), _mo(0) {}

    explicit MoOnTree(
        const m1une::graph::Graph<T>& graph,
        int root_vertex = 0
    ) : root(-1), _n(0), _mo(0) {
        build(graph, root_vertex);
    }

    void build(
        const m1une::graph::Graph<T>& graph,
        int root_vertex = 0
    ) {
        _n = graph.size();
        root = _n == 0 ? -1 : root_vertex;
        entry.assign(_n, -1);
        exit.assign(_n, -1);
        tour.clear();
        tour.reserve(2 * _n);
        _queries.clear();
        _mo = m1une::algo::Mo(2 * _n);
        _hld.build(graph, root_vertex);
        if (_n == 0) return;

        assert(0 <= root && root < _n);
        for (int vertex = 0; vertex < _n; ++vertex) {
            assert(_hld.parent[vertex] != -2);
        }

        std::vector<std::vector<int>> children(_n);
        for (int vertex = 0; vertex < _n; ++vertex) {
            int parent = _hld.parent[vertex];
            if (parent != -1) children[parent].push_back(vertex);
        }

        struct Event {
            int vertex;
            bool leaving;
        };
        std::vector<Event> stack;
        stack.reserve(2 * _n);
        stack.push_back(Event{root, false});
        while (!stack.empty()) {
            Event event = stack.back();
            stack.pop_back();
            int vertex = event.vertex;
            if (event.leaving) {
                exit[vertex] = int(tour.size());
                tour.push_back(vertex);
                continue;
            }

            entry[vertex] = int(tour.size());
            tour.push_back(vertex);
            stack.push_back(Event{vertex, true});
            const auto& child_list = children[vertex];
            for (int index = int(child_list.size()) - 1; index >= 0; --index) {
                stack.push_back(Event{child_list[index], false});
            }
        }
        assert(int(tour.size()) == 2 * _n);
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    int query_count() const {
        return int(_queries.size());
    }

    const std::vector<Query>& queries() const {
        return _queries;
    }

    int parent(int vertex) const {
        check_vertex(vertex);
        return _hld.parent[vertex];
    }

    int parent_edge(int vertex) const {
        check_vertex(vertex);
        return _hld.parent_edge[vertex];
    }

    int depth(int vertex) const {
        check_vertex(vertex);
        return _hld.depth[vertex];
    }

    int lca(int first, int second) const {
        check_vertex(first);
        check_vertex(second);
        return _hld.lca(first, second);
    }

    void reserve(int query_capacity) {
        assert(0 <= query_capacity);
        _queries.reserve(query_capacity);
        _mo.reserve(query_capacity);
    }

    void clear() {
        _queries.clear();
        _mo.clear();
    }

    // Adds an inclusive vertex-path query and returns its insertion-order ID.
    // Vertex and edge queries cannot be mixed in one collection.
    int add_query(int from, int to) {
        return add_path_query(from, to, false);
    }

    // Adds an edge-path query. Each edge is represented by its child vertex.
    int add_edge_query(int from, int to) {
        return add_path_query(from, to, true);
    }

    std::vector<int> order(int block_size = 0) const {
        return _mo.order(block_size);
    }

    // `add(v)` and `remove(v)` maintain the current path. In edge mode, v
    // always represents the real edge parent_edge(v).
    template <class Add, class Remove, class Answer>
    void run(
        Add add,
        Remove remove,
        Answer answer,
        int block_size = 0
    ) const {
        bool edge_mode = !_queries.empty() && _queries.front().edge;
        std::vector<char> active(_n, false);
        auto toggle = [&](int tour_index) {
            int vertex = tour[tour_index];
            if (!edge_mode || vertex != root) {
                if (active[vertex]) {
                    remove(vertex);
                } else {
                    add(vertex);
                }
            }
            active[vertex] = !active[vertex];
        };

        _mo.run(
            toggle,
            toggle,
            [&](int query_id) {
                int extra = _queries[query_id].extra;
                if (extra != -1) {
                    assert(!active[extra]);
                    add(extra);
                }
                answer(query_id);
                if (extra != -1) remove(extra);
            },
            block_size
        );
    }
};

}  // namespace tree
}  // namespace m1une

#endif  // M1UNE_TREE_MO_ON_TREE_HPP
#line 1 "graph/tree/mo_on_tree.hpp"



#include <algorithm>
#include <cassert>
#include <vector>

#line 1 "algo/offline/mo.hpp"



#line 6 "algo/offline/mo.hpp"
#include <cmath>
#include <numeric>
#line 9 "algo/offline/mo.hpp"

namespace m1une {
namespace algo {

// Offline Mo's algorithm for half-open array ranges.
struct Mo {
    struct Query {
        int left;
        int right;
        int id;
    };

   private:
    int _n;
    std::vector<Query> _queries;

   public:
    Mo() : _n(0) {}

    explicit Mo(int n) : _n(n) {
        assert(0 <= n);
    }

    int size() const {
        return _n;
    }

    int query_count() const {
        return int(_queries.size());
    }

    bool empty() const {
        return _queries.empty();
    }

    const std::vector<Query>& queries() const {
        return _queries;
    }

    void reserve(int query_capacity) {
        assert(0 <= query_capacity);
        _queries.reserve(query_capacity);
    }

    void clear() {
        _queries.clear();
    }

    // Adds [left, right) and returns its insertion-order ID.
    int add_query(int left, int right) {
        assert(0 <= left && left <= right && right <= _n);
        int id = query_count();
        _queries.push_back(Query{left, right, id});
        return id;
    }

    // Returns query IDs in Mo order. A non-positive block size selects one
    // automatically.
    std::vector<int> order(int block_size = 0) const {
        int query_size = query_count();
        std::vector<int> result(query_size);
        std::iota(result.begin(), result.end(), 0);
        if (query_size == 0) return result;

        if (block_size <= 0) {
            block_size = std::max(1, int(_n / std::sqrt(static_cast<double>(query_size))));
        }

        std::sort(result.begin(), result.end(), [&](int first, int second) {
            const Query& a = _queries[first];
            const Query& b = _queries[second];
            int first_block = a.left / block_size;
            int second_block = b.left / block_size;
            if (first_block != second_block) {
                return first_block < second_block;
            }
            if (first_block & 1) return a.right > b.right;
            return a.right < b.right;
        });
        return result;
    }

    // Maintains [left, right). Each movement callback receives the array index
    // being inserted or erased. `answer(query_id)` stores or reports a result.
    template <class AddLeft, class AddRight, class RemoveLeft, class RemoveRight, class Answer>
    void run(AddLeft add_left, AddRight add_right, RemoveLeft remove_left, RemoveRight remove_right, Answer answer,
             int block_size = 0) const {
        int left = 0;
        int right = 0;
        for (int query_index : order(block_size)) {
            const Query& query = _queries[query_index];
            while (query.left < left) add_left(--left);
            while (right < query.right) add_right(right++);
            while (left < query.left) remove_left(left++);
            while (query.right < right) remove_right(--right);
            answer(query.id);
        }
    }

    // Convenience overload for statistics whose update is independent of
    // which side moves.
    template <class Add, class Remove, class Answer>
    void run(Add add, Remove remove, Answer answer, int block_size = 0) const {
        run(add, add, remove, remove, answer, block_size);
    }
};

}  // namespace algo
}  // namespace m1une


#line 1 "graph/graph.hpp"



#include <array>
#line 6 "graph/graph.hpp"
#include <utility>
#line 8 "graph/graph.hpp"

namespace m1une {
namespace graph {

template <class T = int>
struct Edge {
    using cost_type = T;

    int from;
    int to;
    T cost;
    int id;
    bool alive;

    Edge() : from(-1), to(-1), cost(T()), id(-1), alive(true) {}
    Edge(int from_, int to_, T cost_ = T(1), int id_ = -1, bool alive_ = true)
        : from(from_), to(to_), cost(cost_), id(id_), alive(alive_) {}

    int other(int v) const {
        assert(v == from || v == to);
        return from ^ to ^ v;
    }
};

template <class T = int>
struct Graph {
    using edge_type = Edge<T>;
    using cost_type = T;

   private:
    struct EdgePositions {
        std::array<std::pair<int, int>, 2> value{};
        int size = 0;

        void push_back(std::pair<int, int> position) {
            assert(size < 2);
            value[size++] = position;
        }
    };

    int _n;
    int _edge_count;
    std::vector<std::vector<edge_type>> _g;
    std::vector<EdgePositions> _edge_positions;

   public:
    Graph() : _n(0), _edge_count(0) {}
    explicit Graph(int n) : _n(n), _edge_count(0), _g(n) {
        assert(0 <= n);
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    int edge_count() const {
        return _edge_count;
    }

    int add_vertex() {
        _g.emplace_back();
        return _n++;
    }

    int add_directed_edge(int from, int to, T cost = T(1)) {
        assert(0 <= from && from < _n);
        assert(0 <= to && to < _n);
        int id = _edge_count++;
        int idx = int(_g[from].size());
        _g[from].push_back(edge_type(from, to, cost, id));
        _edge_positions.emplace_back();
        _edge_positions.back().push_back({from, idx});
        return id;
    }

    int add_edge(int u, int v, T cost = T(1)) {
        assert(0 <= u && u < _n);
        assert(0 <= v && v < _n);
        int id = _edge_count++;
        int u_idx = int(_g[u].size());
        _g[u].push_back(edge_type(u, v, cost, id));
        int v_idx = int(_g[v].size());
        _g[v].push_back(edge_type(v, u, cost, id));
        _edge_positions.emplace_back();
        _edge_positions.back().push_back({u, u_idx});
        _edge_positions.back().push_back({v, v_idx});
        return id;
    }

    void set_edge_alive(int id, bool alive) {
        assert(0 <= id && id < _edge_count);
        for (int i = 0; i < _edge_positions[id].size; ++i) {
            auto [v, idx] = _edge_positions[id].value[i];
            _g[v][idx].alive = alive;
        }
    }

    void erase_edge(int id) {
        set_edge_alive(id, false);
    }

    void revive_edge(int id) {
        set_edge_alive(id, true);
    }

    bool is_edge_alive(int id) const {
        assert(0 <= id && id < _edge_count);
        assert(_edge_positions[id].size != 0);
        auto [v, idx] = _edge_positions[id].value[0];
        return _g[v][idx].alive;
    }

    const std::vector<edge_type>& operator[](int v) const {
        assert(0 <= v && v < _n);
        return _g[v];
    }

    std::vector<edge_type>& operator[](int v) {
        assert(0 <= v && v < _n);
        return _g[v];
    }

    const std::vector<std::vector<edge_type>>& adjacency() const {
        return _g;
    }

    std::vector<std::vector<edge_type>>& adjacency() {
        return _g;
    }

    std::vector<edge_type> edges(bool include_inactive = false) const {
        std::vector<edge_type> result;
        result.reserve(_edge_count);
        std::vector<char> used(_edge_count, false);
        for (int v = 0; v < _n; v++) {
            for (const auto& e : _g[v]) {
                if (!include_inactive && !e.alive) continue;
                if (0 <= e.id && e.id < _edge_count) {
                    if (used[e.id]) continue;
                    used[e.id] = true;
                }
                result.push_back(e);
            }
        }
        return result;
    }

    Graph reversed() const {
        Graph result(_n);
        result._edge_count = _edge_count;
        result._edge_positions.assign(_edge_count, {});
        for (int v = 0; v < _n; v++) {
            for (const auto& e : _g[v]) {
                int idx = int(result._g[e.to].size());
                result._g[e.to].push_back(edge_type(e.to, e.from, e.cost, e.id, e.alive));
                if (0 <= e.id && e.id < _edge_count) result._edge_positions[e.id].push_back({e.to, idx});
            }
        }
        return result;
    }
};

}  // namespace graph
}  // namespace m1une


#line 1 "graph/tree/heavy_light_decomposition.hpp"



#line 8 "graph/tree/heavy_light_decomposition.hpp"

#line 10 "graph/tree/heavy_light_decomposition.hpp"

namespace m1une {
namespace tree {

struct HldPathSegment {
    int l;
    int r;
    bool reversed;
};

template <class T = int>
struct HeavyLightDecomposition {
    using cost_type = T;
    using edge_type = m1une::graph::Edge<T>;

    int root;
    std::vector<int> parent;
    std::vector<int> parent_edge;
    std::vector<int> depth;
    std::vector<T> dist;
    std::vector<int> subtree_size;
    std::vector<int> heavy;
    std::vector<int> head;
    std::vector<int> tin;
    std::vector<int> tout;
    std::vector<int> order;

   private:
    int _n;

    void check_vertex(int v) const {
        assert(0 <= v && v < _n);
        assert(tin[v] != -1);
    }

    static void add_segment(std::vector<HldPathSegment>& result, int l, int r, bool reversed) {
        if (l < r) result.push_back({l, r, reversed});
    }

   public:
    HeavyLightDecomposition() : root(-1), _n(0) {}
    explicit HeavyLightDecomposition(const m1une::graph::Graph<T>& g, int root_ = 0) {
        build(g, root_);
    }

    void build(const m1une::graph::Graph<T>& g, int root_ = 0) {
        _n = g.size();
        root = _n == 0 ? -1 : root_;
        parent.assign(_n, -2);
        parent_edge.assign(_n, -1);
        depth.assign(_n, 0);
        dist.assign(_n, T(0));
        subtree_size.assign(_n, 1);
        heavy.assign(_n, -1);
        head.assign(_n, -1);
        tin.assign(_n, -1);
        tout.assign(_n, -1);
        order.clear();
        order.reserve(_n);
        if (_n == 0) return;
        assert(0 <= root && root < _n);

        std::vector<int> dfs_order;
        dfs_order.reserve(_n);
        std::vector<int> stack = {root};
        parent[root] = -1;
        while (!stack.empty()) {
            int v = stack.back();
            stack.pop_back();
            dfs_order.push_back(v);
            for (const auto& e : g[v]) {
                if (!e.alive) continue;
                if (parent[e.to] != -2) continue;
                parent[e.to] = v;
                parent_edge[e.to] = e.id;
                depth[e.to] = depth[v] + 1;
                dist[e.to] = dist[v] + e.cost;
                stack.push_back(e.to);
            }
        }

        for (int i = int(dfs_order.size()) - 1; i >= 0; i--) {
            int v = dfs_order[i];
            if (parent[v] == -1) continue;
            int p = parent[v];
            subtree_size[p] += subtree_size[v];
            if (heavy[p] == -1 || subtree_size[heavy[p]] < subtree_size[v]) heavy[p] = v;
        }

        order.assign(dfs_order.size(), -1);
        int timer = 0;
        std::vector<std::pair<int, int>> starts = {std::pair<int, int>{root, root}};
        while (!starts.empty()) {
            auto [start, h] = starts.back();
            starts.pop_back();
            for (int v = start; v != -1; v = heavy[v]) {
                head[v] = h;
                tin[v] = timer;
                order[timer++] = v;
                for (auto it = g[v].rbegin(); it != g[v].rend(); ++it) {
                    if (!it->alive) continue;
                    int to = it->to;
                    if (parent[to] != v || to == heavy[v]) continue;
                    starts.push_back({to, to});
                }
            }
        }
        for (int i = int(dfs_order.size()) - 1; i >= 0; i--) {
            int v = dfs_order[i];
            tout[v] = tin[v] + subtree_size[v];
        }
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    bool is_ancestor(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        return tin[u] <= tin[v] && tout[v] <= tout[u];
    }

    int lca(int u, int v) const {
        check_vertex(u);
        check_vertex(v);
        while (head[u] != head[v]) {
            if (depth[head[u]] < depth[head[v]]) std::swap(u, v);
            u = parent[head[u]];
        }
        return depth[u] < depth[v] ? u : v;
    }

    int dist_edges(int u, int v) const {
        int w = lca(u, v);
        return depth[u] + depth[v] - 2 * depth[w];
    }

    T dist_cost(int u, int v) const {
        int w = lca(u, v);
        return dist[u] + dist[v] - dist[w] - dist[w];
    }

    int kth_ancestor(int v, int k) const {
        check_vertex(v);
        assert(0 <= k);
        while (v != -1) {
            int h = head[v];
            int len = depth[v] - depth[h];
            if (k <= len) return order[tin[v] - k];
            k -= len + 1;
            v = parent[h];
        }
        return -1;
    }

    int jump(int from, int to, int k) const {
        check_vertex(from);
        check_vertex(to);
        assert(0 <= k);
        int w = lca(from, to);
        int up_len = depth[from] - depth[w];
        int down_len = depth[to] - depth[w];
        if (up_len + down_len < k) return -1;
        if (k <= up_len) return kth_ancestor(from, k);
        return kth_ancestor(to, down_len - (k - up_len));
    }

    std::pair<int, int> subtree_range(int v, bool edge = false) const {
        check_vertex(v);
        return {tin[v] + (edge ? 1 : 0), tout[v]};
    }

    std::vector<HldPathSegment> path_segments(int u, int v, bool edge = false) const {
        check_vertex(u);
        check_vertex(v);
        std::vector<HldPathSegment> result, down;
        while (head[u] != head[v]) {
            if (depth[head[u]] >= depth[head[v]]) {
                add_segment(result, tin[head[u]], tin[u] + 1, true);
                u = parent[head[u]];
            } else {
                add_segment(down, tin[head[v]], tin[v] + 1, false);
                v = parent[head[v]];
            }
        }

        if (depth[u] >= depth[v]) {
            add_segment(result, tin[v] + (edge ? 1 : 0), tin[u] + 1, true);
        } else {
            add_segment(down, tin[u] + (edge ? 1 : 0), tin[v] + 1, false);
        }
        std::reverse(down.begin(), down.end());
        result.insert(result.end(), down.begin(), down.end());
        return result;
    }

    template <class F>
    void for_each_path(int u, int v, F f, bool edge = false) const {
        for (auto seg : path_segments(u, v, edge)) f(seg.l, seg.r, seg.reversed);
    }
};

}  // namespace tree
}  // namespace m1une


#line 11 "graph/tree/mo_on_tree.hpp"

namespace m1une {
namespace tree {

// Offline Mo's algorithm for static paths in a tree.
template <class T = int>
struct MoOnTree {
    struct Query {
        int from;
        int to;
        int left;
        int right;
        int extra;
        int id;
        bool edge;
    };

    int root;
    std::vector<int> entry;
    std::vector<int> exit;
    std::vector<int> tour;

   private:
    int _n;
    HeavyLightDecomposition<T> _hld;
    m1une::algo::Mo _mo;
    std::vector<Query> _queries;

    void check_vertex(int vertex) const {
        assert(0 <= vertex && vertex < _n);
        assert(entry[vertex] != -1);
    }

    int add_path_query(int from, int to, bool edge) {
        check_vertex(from);
        check_vertex(to);
        assert(_queries.empty() || _queries.front().edge == edge);
        int original_from = from;
        int original_to = to;
        if (entry[from] > entry[to]) std::swap(from, to);

        int ancestor = _hld.lca(from, to);
        int left;
        int right = entry[to] + 1;
        int extra = -1;
        if (ancestor == from) {
            left = entry[from] + int(edge);
        } else {
            left = exit[from];
            if (!edge) extra = ancestor;
        }

        int id = _mo.add_query(left, right);
        _queries.push_back(Query{
            original_from,
            original_to,
            left,
            right,
            extra,
            id,
            edge,
        });
        return id;
    }

   public:
    MoOnTree() : root(-1), _n(0), _mo(0) {}

    explicit MoOnTree(
        const m1une::graph::Graph<T>& graph,
        int root_vertex = 0
    ) : root(-1), _n(0), _mo(0) {
        build(graph, root_vertex);
    }

    void build(
        const m1une::graph::Graph<T>& graph,
        int root_vertex = 0
    ) {
        _n = graph.size();
        root = _n == 0 ? -1 : root_vertex;
        entry.assign(_n, -1);
        exit.assign(_n, -1);
        tour.clear();
        tour.reserve(2 * _n);
        _queries.clear();
        _mo = m1une::algo::Mo(2 * _n);
        _hld.build(graph, root_vertex);
        if (_n == 0) return;

        assert(0 <= root && root < _n);
        for (int vertex = 0; vertex < _n; ++vertex) {
            assert(_hld.parent[vertex] != -2);
        }

        std::vector<std::vector<int>> children(_n);
        for (int vertex = 0; vertex < _n; ++vertex) {
            int parent = _hld.parent[vertex];
            if (parent != -1) children[parent].push_back(vertex);
        }

        struct Event {
            int vertex;
            bool leaving;
        };
        std::vector<Event> stack;
        stack.reserve(2 * _n);
        stack.push_back(Event{root, false});
        while (!stack.empty()) {
            Event event = stack.back();
            stack.pop_back();
            int vertex = event.vertex;
            if (event.leaving) {
                exit[vertex] = int(tour.size());
                tour.push_back(vertex);
                continue;
            }

            entry[vertex] = int(tour.size());
            tour.push_back(vertex);
            stack.push_back(Event{vertex, true});
            const auto& child_list = children[vertex];
            for (int index = int(child_list.size()) - 1; index >= 0; --index) {
                stack.push_back(Event{child_list[index], false});
            }
        }
        assert(int(tour.size()) == 2 * _n);
    }

    int size() const {
        return _n;
    }

    bool empty() const {
        return _n == 0;
    }

    int query_count() const {
        return int(_queries.size());
    }

    const std::vector<Query>& queries() const {
        return _queries;
    }

    int parent(int vertex) const {
        check_vertex(vertex);
        return _hld.parent[vertex];
    }

    int parent_edge(int vertex) const {
        check_vertex(vertex);
        return _hld.parent_edge[vertex];
    }

    int depth(int vertex) const {
        check_vertex(vertex);
        return _hld.depth[vertex];
    }

    int lca(int first, int second) const {
        check_vertex(first);
        check_vertex(second);
        return _hld.lca(first, second);
    }

    void reserve(int query_capacity) {
        assert(0 <= query_capacity);
        _queries.reserve(query_capacity);
        _mo.reserve(query_capacity);
    }

    void clear() {
        _queries.clear();
        _mo.clear();
    }

    // Adds an inclusive vertex-path query and returns its insertion-order ID.
    // Vertex and edge queries cannot be mixed in one collection.
    int add_query(int from, int to) {
        return add_path_query(from, to, false);
    }

    // Adds an edge-path query. Each edge is represented by its child vertex.
    int add_edge_query(int from, int to) {
        return add_path_query(from, to, true);
    }

    std::vector<int> order(int block_size = 0) const {
        return _mo.order(block_size);
    }

    // `add(v)` and `remove(v)` maintain the current path. In edge mode, v
    // always represents the real edge parent_edge(v).
    template <class Add, class Remove, class Answer>
    void run(
        Add add,
        Remove remove,
        Answer answer,
        int block_size = 0
    ) const {
        bool edge_mode = !_queries.empty() && _queries.front().edge;
        std::vector<char> active(_n, false);
        auto toggle = [&](int tour_index) {
            int vertex = tour[tour_index];
            if (!edge_mode || vertex != root) {
                if (active[vertex]) {
                    remove(vertex);
                } else {
                    add(vertex);
                }
            }
            active[vertex] = !active[vertex];
        };

        _mo.run(
            toggle,
            toggle,
            [&](int query_id) {
                int extra = _queries[query_id].extra;
                if (extra != -1) {
                    assert(!active[extra]);
                    add(extra);
                }
                answer(query_id);
                if (extra != -1) remove(extra);
            },
            block_size
        );
    }
};

}  // namespace tree
}  // namespace m1une
Back to top page