graph/tree/auxiliary_tree.hpp

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

Depends on

• graph/tree/heavy_light_decomposition.hpp

Code

#ifndef KK2_GRAPH_TREE_AUXILIARY_TREE_HPP
#define KK2_GRAPH_TREE_AUXILIARY_TREE_HPP 1

#include <algorithm>
#include <cassert>
#include <functional>
#include <stack>
#include <utility>
#include <vector>

#include "heavy_light_decomposition.hpp"

namespace kk2 {

template <typename G> struct AuxiliaryTree {
    static_assert(!G::directed, "AuxiliaryTree requires undirected graph");

    G g;
    HeavyLightDecomposition<G> hld;

    AuxiliaryTree(const G &g_, int root_ = 0) : g(g_), hld(g, root_) {}

    std::pair<std::vector<std::vector<int>>, std::vector<int>> get(std::vector<int> ps) {
        if (ps.empty()) return {};
        std::sort(
            std::begin(ps), std::end(ps), [&](int i, int j) { return hld.in[i] < hld.in[j]; });
        for (int i = 0, ps_size = ps.size(); i < ps_size - 1; i++) {
            ps.push_back(hld.lca(ps[i], ps[i + 1]));
        }
        std::sort(
            std::begin(ps), std::end(ps), [&](int i, int j) { return hld.in[i] < hld.in[j]; });
        ps.erase(std::unique(std::begin(ps), std::end(ps)), std::end(ps));

        std::vector<std::vector<int>> aux(ps.size());
        std::stack<int> st;
        st.emplace(0);
        for (int i = 1; i < (int)ps.size(); i++) {
            int l = hld.lca(ps[i], ps[st.top()]);
            while (ps[st.top()] != l) st.pop();
            aux[st.top()].emplace_back(i);
            st.emplace(i);
        }
        return std::make_pair(aux, ps);
    }
};

} // namespace kk2

#endif // KK2_GRAPH_TREE_AUXILIARY_TREE_HPP

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



#include <algorithm>
#include <cassert>
#include <functional>
#include <stack>
#include <utility>
#include <vector>

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



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

namespace kk2 {

template <typename G> struct HeavyLightDecomposition {
    static_assert(!G::directed, "HeavyLightDecomposition requires undirected graph");

    G &g;
    int root, id;
    std::vector<int> sz, in, out, head, par, dep, edge_idx;

    // e.idはedgesのindexでないといけない
    HeavyLightDecomposition(G &g_, int root_ = 0)
        : g(g_),
          root(root_),
          id(0),
          sz(g.size(), 0),
          in(g.size(), -1),
          out(g.size(), -1),
          head(g.size(), root),
          par(g.size(), root),
          dep(g.size(), 0),
          edge_idx(g.size() - 1, -1) {
        init();
    }

    int get_edge_idx(int i) const { return edge_idx[i]; }

    std::pair<int, int> get_node_idx(int u) const { return std::make_pair(in[u], out[u]); }

    template <typename F> void path_query(int u, int v, bool is_node_query, const F &f) {
        int l = lca(u, v);
        for (auto &[a, b] : ascend(u, l)) {
            int s = a + 1, t = b;
            s > t ? f(t, s) : f(s, t);
        }
        if (is_node_query) f(in[l], in[l] + 1);
        for (auto &[a, b] : descend(l, v)) {
            int s = a, t = b + 1;
            s > t ? f(t, s) : f(s, t);
        }
    }

    template <typename F>
    void path_noncommutative_query(int u, int v, bool is_node_query, const F &f) {
        int l = lca(u, v);
        for (auto &[a, b] : ascend(u, l)) f(a + 1, b);
        if (is_node_query) f(in[l], in[l] + 1);
        for (auto &[a, b] : descend(l, v)) f(a, b + 1);
    }

    template <typename F> void subtree_query(int u, bool is_vertex_query, const F &f) {
        f(in[u] + (int)!is_vertex_query, out[u]);
    }

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

    int dist(int u, int v) const { return dep[u] + dep[v] - 2 * dep[lca(u, v)]; }

  private:
    void init() {
        auto dfs_sz = [&](auto self, int now) -> void {
            sz[now] = 1;
            for (auto &e : g[now]) {
                if ((int)e == par[now]) {
                    if (g[now].size() >= 2 and e == g[now][0]) std::swap(e, g[now][1]);
                    else continue;
                }
                par[(int)e] = now;
                dep[(int)e] = dep[now] + 1;
                self(self, (int)e);
                sz[now] += sz[(int)e];
                if (sz[(int)e] > sz[(int)g[now][0]]) std::swap(e, g[now][0]);
            }
        };
        dfs_sz(dfs_sz, root);

        auto dfs_hld = [&](auto self, int now) -> void {
            in[now] = id++;
            for (auto &e : g[now]) {
                if ((int)e == par[now]) continue;
                head[(int)e] = ((int)e == (int)g[now][0] ? head[now] : (int)e);
                edge_idx[e.id] = id;
                self(self, (int)e);
            }
            out[now] = id;
        };
        dfs_hld(dfs_hld, root);
    }

    // [u, v)
    std::vector<std::pair<int, int>> ascend(int u, int v) const {
        std::vector<std::pair<int, int>> res;
        while (head[u] != head[v]) {
            res.emplace_back(in[u], in[head[u]]);
            u = par[head[u]];
        }
        if (u != v) res.emplace_back(in[u], in[v] + 1);
        return res;
    }

    // (u, v]
    std::vector<std::pair<int, int>> descend(int u, int v) const {
        if (u == v) return {};
        if (head[u] == head[v]) return {std::make_pair(in[u] + 1, in[v])};
        auto res = descend(u, par[head[v]]);
        res.emplace_back(in[head[v]], in[v]);
        return res;
    }
};

} // namespace kk2


#line 12 "graph/tree/auxiliary_tree.hpp"

namespace kk2 {

template <typename G> struct AuxiliaryTree {
    static_assert(!G::directed, "AuxiliaryTree requires undirected graph");

    G g;
    HeavyLightDecomposition<G> hld;

    AuxiliaryTree(const G &g_, int root_ = 0) : g(g_), hld(g, root_) {}

    std::pair<std::vector<std::vector<int>>, std::vector<int>> get(std::vector<int> ps) {
        if (ps.empty()) return {};
        std::sort(
            std::begin(ps), std::end(ps), [&](int i, int j) { return hld.in[i] < hld.in[j]; });
        for (int i = 0, ps_size = ps.size(); i < ps_size - 1; i++) {
            ps.push_back(hld.lca(ps[i], ps[i + 1]));
        }
        std::sort(
            std::begin(ps), std::end(ps), [&](int i, int j) { return hld.in[i] < hld.in[j]; });
        ps.erase(std::unique(std::begin(ps), std::end(ps)), std::end(ps));

        std::vector<std::vector<int>> aux(ps.size());
        std::stack<int> st;
        st.emplace(0);
        for (int i = 1; i < (int)ps.size(); i++) {
            int l = hld.lca(ps[i], ps[st.top()]);
            while (ps[st.top()] != l) st.pop();
            aux[st.top()].emplace_back(i);
            st.emplace(i);
        }
        return std::make_pair(aux, ps);
    }
};

} // namespace kk2

Back to top page