tree/lca_by_rmq.hpp

• View this file on GitHub
• Last update: 2025-02-04 01:12:01+09:00
• Include: #include "tree/lca_by_rmq.hpp"

Depends on

• Sparse Table (sparse_table.hpp)

Required by

• tree/auxiliary_tree.hpp

Code

#include "../sparse_table.hpp"

class LCASolverRMQ {
   public:
    LCASolverRMQ() = default;

    template <typename Cost>
    LCASolverRMQ(Tree<Cost> &tree) {
        build(tree);
    }

    template <typename Cost>
    void build(Tree<Cost> &tree) {
        if (!tree.built) {
            tree.build();
        }
        int n = tree.n;
        in = tree.in;
        vector<S> s(2 * n);
        if (*max_element(tree.depth.begin(), tree.depth.end()) > INT_MAX) {
            Tree<int> t(n);
            for (int i = 0; i < n; ++i) {
                t.g[i].reserve(tree.g[i].size());
                for (auto [u, c] : tree.g[i]) {
                    t.add_edge(i, u);
                }
            }
            t.build();
            for (int i = 0; i < 2 * n; ++i) {
                s[i].first = t.depth[t.ord_inv[i]];
                unsigned int edge = t.es[i];
                if (edge >> 31 & 1) {
                    s[i].second = t.par[~edge];
                } else {
                    s[i].second = t.par[edge];
                }
            }
        } else {
            for (int i = 0; i < 2 * n; ++i) {
                s[i].first = tree.depth[tree.ord_inv[i]];
                unsigned int edge = tree.es[i];
                if (edge >> 31 & 1) {
                    s[i].second = tree.par[~edge];
                } else {
                    s[i].second = tree.par[edge];
                }
            }
        }
        sp = decltype(sp)(s);
    }

    int lca(int u, int v) {
        if (in[u] > in[v]) swap(u, v);
        return sp.prod(in[u] + 1, in[v] + 1).second;
    }

   private:
    vector<int> in;
    using S = pair<int, int>;
    static S op(S a, S b) { return min(a, b); }
    SparseTable<S, op> sp;
};

#line 1 "sparse_table.hpp"
template <typename S, auto op>
class SparseTable {
    static_assert(std::is_convertible_v<decltype(op), std::function<S(S, S)>>, "op must work as S(S, S)");

   public:
    explicit SparseTable(const vector<S>& v) : n(v.size()), log_table(n + 1) {
        for (int i = 2; i < n + 1; ++i) {
            log_table[i] = log_table[i >> 1] + 1;
        }

        table.resize(log_table[n] + 1, vector<S>(n));

        for (int i = 0; i < n; ++i) {
            table[0][i] = v[i];
        }

        for (int i = 1; (1 << i) <= n; ++i) {
            for (int j = 0; j + (1 << i) <= n; ++j) {
                table[i][j] = op(table[i - 1][j], table[i - 1][j + (1 << (i - 1))]);
            }
        }
    }

    S prod(int l, int r) const {
        assert(0 <= l && l <= r && r <= n);
        const int k = log_table[r - l];

        return op(table[k][l], table[k][r - (1 << k)]);
    }

   private:
    int n;

    vector<int> log_table;
    vector<vector<S>> table;
};
#line 2 "tree/lca_by_rmq.hpp"

class LCASolverRMQ {
   public:
    LCASolverRMQ() = default;

    template <typename Cost>
    LCASolverRMQ(Tree<Cost> &tree) {
        build(tree);
    }

    template <typename Cost>
    void build(Tree<Cost> &tree) {
        if (!tree.built) {
            tree.build();
        }
        int n = tree.n;
        in = tree.in;
        vector<S> s(2 * n);
        if (*max_element(tree.depth.begin(), tree.depth.end()) > INT_MAX) {
            Tree<int> t(n);
            for (int i = 0; i < n; ++i) {
                t.g[i].reserve(tree.g[i].size());
                for (auto [u, c] : tree.g[i]) {
                    t.add_edge(i, u);
                }
            }
            t.build();
            for (int i = 0; i < 2 * n; ++i) {
                s[i].first = t.depth[t.ord_inv[i]];
                unsigned int edge = t.es[i];
                if (edge >> 31 & 1) {
                    s[i].second = t.par[~edge];
                } else {
                    s[i].second = t.par[edge];
                }
            }
        } else {
            for (int i = 0; i < 2 * n; ++i) {
                s[i].first = tree.depth[tree.ord_inv[i]];
                unsigned int edge = tree.es[i];
                if (edge >> 31 & 1) {
                    s[i].second = tree.par[~edge];
                } else {
                    s[i].second = tree.par[edge];
                }
            }
        }
        sp = decltype(sp)(s);
    }

    int lca(int u, int v) {
        if (in[u] > in[v]) swap(u, v);
        return sp.prod(in[u] + 1, in[v] + 1).second;
    }

   private:
    vector<int> in;
    using S = pair<int, int>;
    static S op(S a, S b) { return min(a, b); }
    SparseTable<S, op> sp;
};

Back to top page