Auxiliary Tree
(tree/auxiliary_tree.hpp)

• View this file on GitHub
• Last update: 2024-03-02 20:34:40+09:00
• Include: #include "tree/auxiliary_tree.hpp"

Description

Auxiliary tree は，木の頂点のある部分集合と集合内の頂点間の LCA のみを含み，頂点同士の関係を保ったままもとの木を圧縮したものである．

Operations

• AuxiliaryTree(vector<vector<int>> G, int root)
  • 前計算をする
  • 時間計算量: $O(n\log n)$
• vector<pair<int, int>> query(vector<int> vs)
  • 頂点集合 vs から定まる auxiliary tree の辺を返す
  • 時間計算量: $O(k\log n)$

Reference

• LCAをベースに構築するAuxiliary Treeのメモ
• verify: https://atcoder.jp/contests/typical90/submissions/30186948

Depends on

• Lowest Common Ancestor (tree/lca.hpp)

Code

#pragma once
#include <stack>
#include <utility>
#include <vector>

#include "lca.hpp"

class AuxiliaryTree {
   public:
    AuxiliaryTree() = default;
    AuxiliaryTree(std::vector<std::vector<int>>& G, int root)
        : G(G), lca(G, root), ord(G.size()), depth(G.size()) {
        dfs(root, -1);
    }

    std::vector<std::pair<int, int>> query(std::vector<int> vs) {
        std::ranges::sort(vs, {}, [&](int v) { return ord[v]; });
        std::vector<std::pair<int, int>> edges;
        std::stack<int> st;
        st.push(vs[0]);
        for (int i = 1; i < (int)vs.size(); ++i) {
            int l = lca.query(vs[i - 1], vs[i]);
            if (l != vs[i - 1]) {
                int u = st.top();
                st.pop();
                while (!st.empty() && depth[l] < depth[st.top()]) {
                    edges.push_back({u, st.top()});
                    u = st.top();
                    st.pop();
                }
                edges.push_back({u, l});
                if (st.empty() || st.top() != l) st.push(l);
            }
            st.push(vs[i]);
        }
        while (st.size() > 1) {
            int u = st.top();
            st.pop();
            edges.push_back({u, st.top()});
        }
        return edges;
    }

   private:
    std::vector<std::vector<int>> G;
    LCA lca;
    std::vector<int> ord, depth;
    int k = 0;

    void dfs(int v, int p) {
        ord[v] = k++;
        for (int c : G[v]) {
            if (c != p) {
                depth[c] = depth[v] + 1;
                dfs(c, v);
            }
        }
    }
};

#line 2 "tree/auxiliary_tree.hpp"
#include <stack>
#include <utility>
#include <vector>

#line 2 "tree/lca.hpp"
#include <algorithm>
#include <bit>
#line 5 "tree/lca.hpp"

class LCA {
   public:
    LCA() = default;
    LCA(const std::vector<std::vector<int>>& G, int root)
        : G(G),
          LOG(std::bit_width(G.size())),
          depth(G.size()),
          table(LOG, std::vector<int>(G.size(), -1)) {
        dfs(root, -1, 0);

        for (int k = 0; k < LOG - 1; ++k) {
            for (int v = 0; v < (int)G.size(); ++v) {
                if (table[k][v] >= 0) {
                    table[k + 1][v] = table[k][table[k][v]];
                }
            }
        }
    }

    int query(int u, int v) const {
        if (depth[u] > depth[v]) std::swap(u, v);

        // go up to the same depth
        for (int k = 0; k < LOG; ++k) {
            if ((depth[v] - depth[u]) >> k & 1) {
                v = table[k][v];
            }
        }
        if (u == v) return u;

        for (int k = LOG - 1; k >= 0; --k) {
            if (table[k][u] != table[k][v]) {
                u = table[k][u];
                v = table[k][v];
            }
        }
        return table[0][u];
    }

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

    int parent(int v, int k) const {
        for (int i = LOG - 1; i >= 0; --i) {
            if (k >= (1 << i)) {
                v = table[i][v];
                k -= 1 << i;
            }
        }
        return v;
    }

    int jump(int u, int v, int k) const {
        int l = query(u, v);
        int du = depth[u] - depth[l];
        int dv = depth[v] - depth[l];
        if (du + dv < k) return -1;
        if (k < du) return parent(u, k);
        return parent(v, du + dv - k);
    }

   protected:
    const std::vector<std::vector<int>>& G;
    const int LOG;
    std::vector<int> depth;
    std::vector<std::vector<int>> table;

    void dfs(int v, int p, int d) {
        table[0][v] = p;
        depth[v] = d;
        for (int c : G[v]) {
            if (c != p) dfs(c, v, d + 1);
        }
    }
};
#line 7 "tree/auxiliary_tree.hpp"

class AuxiliaryTree {
   public:
    AuxiliaryTree() = default;
    AuxiliaryTree(std::vector<std::vector<int>>& G, int root)
        : G(G), lca(G, root), ord(G.size()), depth(G.size()) {
        dfs(root, -1);
    }

    std::vector<std::pair<int, int>> query(std::vector<int> vs) {
        std::ranges::sort(vs, {}, [&](int v) { return ord[v]; });
        std::vector<std::pair<int, int>> edges;
        std::stack<int> st;
        st.push(vs[0]);
        for (int i = 1; i < (int)vs.size(); ++i) {
            int l = lca.query(vs[i - 1], vs[i]);
            if (l != vs[i - 1]) {
                int u = st.top();
                st.pop();
                while (!st.empty() && depth[l] < depth[st.top()]) {
                    edges.push_back({u, st.top()});
                    u = st.top();
                    st.pop();
                }
                edges.push_back({u, l});
                if (st.empty() || st.top() != l) st.push(l);
            }
            st.push(vs[i]);
        }
        while (st.size() > 1) {
            int u = st.top();
            st.pop();
            edges.push_back({u, st.top()});
        }
        return edges;
    }

   private:
    std::vector<std::vector<int>> G;
    LCA lca;
    std::vector<int> ord, depth;
    int k = 0;

    void dfs(int v, int p) {
        ord[v] = k++;
        for (int c : G[v]) {
            if (c != p) {
                depth[c] = depth[v] + 1;
                dfs(c, v);
            }
        }
    }
};

Back to top page