最小共通祖先(LCA)
(graph/tree/lca.hpp)

• View this file on GitHub
• Last update: 2023-11-04 18:06:00+09:00
• Include: #include "graph/tree/lca.hpp"

ダブリングによる実装。

HL分解でLCAを求める方が早いような気がするので、これが必要かどうかは不明。

Depends on

• グラフ構造体 (graph/graph_template.hpp)

Verified with

• test/library-checker/tree/lowest-common-ancestor.test.cpp

Code

#pragma once

#include <cmath>
#include <vector>

#include "../graph_template.hpp"

template <typename Cost = int> class LCA {
  public:
    LCA() = default;
    explicit LCA(const Graph<Cost> &G, int root = 0)
        : G(G), root(root), LOG((int)log2(G.size()) + 1), depth(G.size()),
          parent(LOG, std::vector<int>(G.size())) {
        dfs(root, -1, 0);
        for(int k = 0; k + 1 < LOG; k++) {
            for(int i = 0; i < (int)G.size(); i++) {
                if(parent[k][i] < 0) {
                    parent[k + 1][i] = -1;
                } else {
                    parent[k + 1][i] = parent[k][parent[k][i]];
                }
            }
        }
    }
    int get_lca(int u, int v) {
        if(depth[u] > depth[v]) std::swap(u, v);
        for(int k = 0; k < LOG; k++) {
            if((depth[u] - depth[v]) >> k & 1) v = parent[k][v];
        }
        if(u == v) return u;
        for(int k = LOG - 1; k >= 0; k--) {
            if(parent[k][u] != parent[k][v]) {
                u = parent[k][u];
                v = parent[k][v];
            }
        }
        return parent[0][u];
    }
    int get_dist(int u, int v) {
        return (depth[u] + depth[v] - 2 * depth[get_lca(u, v)]);
    }

  private:
    Graph<Cost> G;
    const int root;
    const int LOG;
    std::vector<int> depth;
    std::vector<std::vector<int>> parent;

    void dfs(int u, int p, int d) {
        depth[u] = d;
        parent[0][u] = p;
        for(int v : G[u]) {
            if(v == p) continue;
            dfs(v, u, d + 1);
        }
    }
};

#line 2 "graph/tree/lca.hpp"

#include <cmath>
#include <vector>

#line 2 "graph/graph_template.hpp"

#include <cassert>
#line 5 "graph/graph_template.hpp"

template <typename Cost = int> struct Edge {
    int from, to;
    Cost cost;
    int id;
    Edge() = default;
    explicit Edge(int from, int to, Cost cost = 1, int id = -1)
        : from(from), to(to), cost(cost), id(id) {}
    operator int() const { return to; }
};

template <typename Cost = int> class Graph {
  public:
    Graph() = default;
    explicit Graph(int N) : N(N), M(0), G(N) {}

    inline void add_directed_edge(int from, int to, Cost cost = 1) {
        assert(0 <= from && from < N);
        assert(0 <= to && to < N);
        G[from].emplace_back(from, to, cost, M++);
    }

    inline void add_undirected_edge(int from, int to, Cost cost = 1) {
        assert(0 <= from && from < N);
        assert(0 <= to && to < N);
        G[from].emplace_back(from, to, cost, M);
        G[to].emplace_back(to, from, cost, M++);
    }

    inline size_t size() const { return G.size(); }
    inline std::vector<Edge<Cost>> &operator[](const int &i) { return G[i]; }
    inline const std::vector<Edge<Cost>> &operator[](const int &i) const {
        return G[i];
    }

  protected:
    int N, M;
    std::vector<std::vector<Edge<Cost>>> G;
};

template <class Cost = int> using Edges = std::vector<Edge<Cost>>;
#line 7 "graph/tree/lca.hpp"

template <typename Cost = int> class LCA {
  public:
    LCA() = default;
    explicit LCA(const Graph<Cost> &G, int root = 0)
        : G(G), root(root), LOG((int)log2(G.size()) + 1), depth(G.size()),
          parent(LOG, std::vector<int>(G.size())) {
        dfs(root, -1, 0);
        for(int k = 0; k + 1 < LOG; k++) {
            for(int i = 0; i < (int)G.size(); i++) {
                if(parent[k][i] < 0) {
                    parent[k + 1][i] = -1;
                } else {
                    parent[k + 1][i] = parent[k][parent[k][i]];
                }
            }
        }
    }
    int get_lca(int u, int v) {
        if(depth[u] > depth[v]) std::swap(u, v);
        for(int k = 0; k < LOG; k++) {
            if((depth[u] - depth[v]) >> k & 1) v = parent[k][v];
        }
        if(u == v) return u;
        for(int k = LOG - 1; k >= 0; k--) {
            if(parent[k][u] != parent[k][v]) {
                u = parent[k][u];
                v = parent[k][v];
            }
        }
        return parent[0][u];
    }
    int get_dist(int u, int v) {
        return (depth[u] + depth[v] - 2 * depth[get_lca(u, v)]);
    }

  private:
    Graph<Cost> G;
    const int root;
    const int LOG;
    std::vector<int> depth;
    std::vector<std::vector<int>> parent;

    void dfs(int u, int p, int d) {
        depth[u] = d;
        parent[0][u] = p;
        for(int v : G[u]) {
            if(v == p) continue;
            dfs(v, u, d + 1);
        }
    }
};

Back to top page