tree/auxiliary_tree.hpp

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

Depends on

• Sparse Table (sparse_table.hpp)
• tree/lca_by_rmq.hpp

Code

#include "lca_by_rmq.hpp"

template <typename Cost>
class AuxiliaryTreeSolver {
   public:
    AuxiliaryTreeSolver(Tree<Cost> &tree) { init(tree); }

    void init(Tree<Cost> &tree) {
        if (!tree.built) {
            tree.build();
        }
        n = tree.n;
        in = tree.in;
        index_backet.resize(n, -1);
        depth = tree.depth;
        lca_solver.build(tree);
    }

    pair<vector<tuple<int, int, Cost>>, vector<bool>> aux(vector<int> x) {
        int k = x.size();
        ranges::sort(x, ranges::less{}, [&](int i) { return in[i]; });
        vector<pair<int, int>> edge_list;
        edge_list.reserve(2 * k);
        stack<int> st;
        st.push(x[0]);
        for (int i = 1; i < k; ++i) {
            int w = lca_solver.lca(x[i - 1], x[i]);
            while (!st.empty() && depth[w] <= depth[st.top()]) {
                int u = st.top();
                st.pop();
                if (st.empty() || depth[st.top()] < depth[w]) {
                    if (w != u) {
                        edge_list.emplace_back(w, u);
                    }
                } else {
                    if (st.top() != u) {
                        edge_list.emplace_back(st.top(), u);
                    }
                }
            }
            st.push(w);
            st.push(x[i]);
        }
        int root = -1;
        while (!st.empty()) {
            int u = st.top();
            root = u;
            st.pop();
            if (!st.empty()) {
                if (st.top() != u) edge_list.emplace_back(st.top(), u);
            }
        }

        int k2 = edge_list.size() + 1;
        vector<int> v(k2);
        v[0] = root;
        for (int i = 0; i < k2 - 1; ++i) v[i + 1] = edge_list[i].second;
        for (int i = 0; i < k2; ++i) {
            index_backet[v[i]] = i;
        }
        vector<bool> type(k2);
        for (int i = 0; i < k; ++i) {
            type[index_backet[x[i]]] = true;
        }
        vector<tuple<int, int, Cost>> res(k2 - 1);
        for (int i = 0; i < k2 - 1; ++i) {
            get<0>(res[i]) = index_backet[edge_list[i].first];
            get<1>(res[i]) = index_backet[edge_list[i].second];
            get<2>(res[i]) = depth[edge_list[i].second] - depth[edge_list[i].first];
        }
        return {res, type};
    }

    int n;

   private:
    vector<int> in;
    vector<Cost> depth;
    LCASolverRMQ lca_solver;
    vector<int> index_backet;
};

#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;
};
#line 2 "tree/auxiliary_tree.hpp"

template <typename Cost>
class AuxiliaryTreeSolver {
   public:
    AuxiliaryTreeSolver(Tree<Cost> &tree) { init(tree); }

    void init(Tree<Cost> &tree) {
        if (!tree.built) {
            tree.build();
        }
        n = tree.n;
        in = tree.in;
        index_backet.resize(n, -1);
        depth = tree.depth;
        lca_solver.build(tree);
    }

    pair<vector<tuple<int, int, Cost>>, vector<bool>> aux(vector<int> x) {
        int k = x.size();
        ranges::sort(x, ranges::less{}, [&](int i) { return in[i]; });
        vector<pair<int, int>> edge_list;
        edge_list.reserve(2 * k);
        stack<int> st;
        st.push(x[0]);
        for (int i = 1; i < k; ++i) {
            int w = lca_solver.lca(x[i - 1], x[i]);
            while (!st.empty() && depth[w] <= depth[st.top()]) {
                int u = st.top();
                st.pop();
                if (st.empty() || depth[st.top()] < depth[w]) {
                    if (w != u) {
                        edge_list.emplace_back(w, u);
                    }
                } else {
                    if (st.top() != u) {
                        edge_list.emplace_back(st.top(), u);
                    }
                }
            }
            st.push(w);
            st.push(x[i]);
        }
        int root = -1;
        while (!st.empty()) {
            int u = st.top();
            root = u;
            st.pop();
            if (!st.empty()) {
                if (st.top() != u) edge_list.emplace_back(st.top(), u);
            }
        }

        int k2 = edge_list.size() + 1;
        vector<int> v(k2);
        v[0] = root;
        for (int i = 0; i < k2 - 1; ++i) v[i + 1] = edge_list[i].second;
        for (int i = 0; i < k2; ++i) {
            index_backet[v[i]] = i;
        }
        vector<bool> type(k2);
        for (int i = 0; i < k; ++i) {
            type[index_backet[x[i]]] = true;
        }
        vector<tuple<int, int, Cost>> res(k2 - 1);
        for (int i = 0; i < k2 - 1; ++i) {
            get<0>(res[i]) = index_backet[edge_list[i].first];
            get<1>(res[i]) = index_backet[edge_list[i].second];
            get<2>(res[i]) = depth[edge_list[i].second] - depth[edge_list[i].first];
        }
        return {res, type};
    }

    int n;

   private:
    vector<int> in;
    vector<Cost> depth;
    LCASolverRMQ lca_solver;
    vector<int> index_backet;
};

Back to top page