graph/lowest_common_ancestor.yosupo.test.cpp

View this file on GitHub
Last update: 2021-08-30 04:35:37+09:00
Problem: https://judge.yosupo.jp/problem/lca

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#include "../graph/lowest_common_ancestor.hpp"

#include "../utils/macros.hpp"
#include "../hack/stack_pivot.hpp"
#include <cstdio>
#include <vector>
using namespace std;

int moin() {
    // read a tree
    int n, q; scanf("%d%d", &n, &q);
    vector<vector<int> > g(n);
    REP (i, n - 1) {
        int p; scanf("%d", &p);
        g[i + 1].push_back(p);
        g[p].push_back(i + 1);
    }

    // construct the LCA
    constexpr int root = 0;
    lowest_common_ancestor lca(root, g);

    // answer to queries
    while (q --) {
        int u, v; scanf("%d%d", &u, &v);
        printf("%d\n", lca(u, v));
    }
    return 0;
}

STACK_PIVOT_MAIN(moin)

#line 1 "graph/lowest_common_ancestor.yosupo.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/lca"
#line 2 "graph/lowest_common_ancestor.hpp"
#include <algorithm>
#include <cassert>
#include <functional>
#include <utility>
#include <vector>
#line 2 "utils/macros.hpp"
#define REP(i, n) for (int i = 0; (i) < (int)(n); ++ (i))
#define REP3(i, m, n) for (int i = (m); (i) < (int)(n); ++ (i))
#define REP_R(i, n) for (int i = (int)(n) - 1; (i) >= 0; -- (i))
#define REP3R(i, m, n) for (int i = (int)(n) - 1; (i) >= (int)(m); -- (i))
#define ALL(x) std::begin(x), std::end(x)
#line 5 "data_structure/sparse_table.hpp"

/**
 * @brief Sparse Table (idempotent monoid)
 * @note the unit is required just for convenience
 * @note $O(N \log N)$ space
 */
template <class IdempotentMonoid>
struct sparse_table {
    typedef typename IdempotentMonoid::value_type value_type;
    std::vector<std::vector<value_type> > table;
    IdempotentMonoid mon;
    sparse_table() = default;

    /**
     * @note $O(N \log N)$ time
     */
    template <class InputIterator>
    sparse_table(InputIterator first, InputIterator last, const IdempotentMonoid & mon_ = IdempotentMonoid())
            : mon(mon_) {
        table.emplace_back(first, last);
        int n = table[0].size();
        int log_n = 32 - __builtin_clz(n);
        table.resize(log_n, std::vector<value_type>(n));
        REP (k, log_n - 1) {
            REP (i, n) {
                table[k + 1][i] = i + (1ll << k) < n ?
                    mon.mult(table[k][i], table[k][i + (1ll << k)]) :
                    table[k][i];
            }
        }
    }

    /**
     * @note $O(1)$
     */
    value_type range_get(int l, int r) const {
        if (l == r) return mon.unit();  // if there is no unit, remove this line
        assert (0 <= l and l < r and r <= (int)table[0].size());
        int k = 31 - __builtin_clz(r - l);  // log2
        return mon.mult(table[k][l], table[k][r - (1ll << k)]);
    }
};
#line 3 "monoids/min_index.hpp"
#include <climits>
#include <limits>
#line 6 "monoids/min_index.hpp"

/**
 * @note a semilattice
 */
template <class T>
struct min_index_monoid {
    typedef std::pair<T, int> value_type;
    value_type unit() const { return std::make_pair(std::numeric_limits<T>::max(), INT_MAX); }
    value_type mult(value_type a, value_type b) const { return std::min(a, b); }
};
#line 9 "graph/lowest_common_ancestor.hpp"

/**
 * @brief lowest common ancestor / 最小共通祖先 (前処理 $O(N)$ + $O(1)$, $\pm 1$ RMQ and sparse table)
 * @see https://www.slideshare.net/yumainoue965/lca-and-rmq
 * @note verified http://www.utpc.jp/2011/problems/travel.html
 */
struct lowest_common_ancestor {
    sparse_table<min_index_monoid<int> > table;
    std::vector<int> index;
    lowest_common_ancestor() = default;
    /**
     * @note $O(N)$
     * @param g is an adjacent list of a tree
     * @note you can easily modify this to accept forests
     */
    lowest_common_ancestor(int root, std::vector<std::vector<int> > const & g) {
        std::vector<std::pair<int, int> > tour;
        index.assign(g.size(), -1);
        dfs(root, -1, 0, g, tour);
        table = sparse_table<min_index_monoid<int> >(ALL(tour));
    }
private:
    /**
     * @note sometimes causes stack overflow without ulimit -s unlimited
     */
    void dfs(int x, int parent, int depth, std::vector<std::vector<int> > const & g, std::vector<std::pair<int, int> > & tour) {
        index[x] = tour.size();
        tour.emplace_back(depth, x);
        for (int y : g[x]) if (y != parent) {
            dfs(y, x, depth + 1, g, tour);
            tour.emplace_back(depth, x);
        }
    }
public:
    /**
     * @note $O(1)$
     */
    int operator () (int x, int y) const {
        assert (0 <= x and x < index.size());
        assert (0 <= y and y < index.size());
        x = index[x];
        y = index[y];
        if (x > y) std::swap(x, y);
        return table.range_get(x, y + 1).second;
    }
    int get_depth(int x) const {
        assert (0 <= x and x < index.size());
        return table.range_get(index[x], index[x] + 1).first;
    }
    int get_dist(int x, int y) const {
        assert (0 <= x and x < index.size());
        assert (0 <= y and y < index.size());
        int z = (*this)(x, y);
        return get_depth(x) + get_depth(y) - 2 * get_depth(z);
    }
};
#line 3 "graph/lowest_common_ancestor.yosupo.test.cpp"

#line 2 "hack/stack_pivot.hpp"
#include <cstdlib>

/**
 * @note This is a workaround for WSL. We cannot use ulimit -s unlimited on the environment.
 * @note To use such techniques, you should take care of the alignment of rsp. If not, you'll get SIGSEGV around XMM registers.
 */

#define BEGIN_STACK_PIVOT(STACK_SIZE) \
    static volatile char *old_stack; \
    asm volatile("mov %%rsp, %0" : "=r" (old_stack) ); \
    char *new_stack = ((char *)malloc(STACK_SIZE) + (STACK_SIZE) - 0x10); \
    asm volatile("mov %0, %%rsp" : : "r" (new_stack) );

#define END_STACK_PIVOT() \
    asm volatile("mov %0, %%rsp" : : "r" (old_stack) );

#define STACK_PIVOT_MAIN(moin) \
    int main() { \
        BEGIN_STACK_PIVOT(1 << 28) \
        static int returncode = moin(); \
        END_STACK_PIVOT() \
        return returncode; \
    }
#line 6 "graph/lowest_common_ancestor.yosupo.test.cpp"
#include <cstdio>
#line 8 "graph/lowest_common_ancestor.yosupo.test.cpp"
using namespace std;

int moin() {
    // read a tree
    int n, q; scanf("%d%d", &n, &q);
    vector<vector<int> > g(n);
    REP (i, n - 1) {
        int p; scanf("%d", &p);
        g[i + 1].push_back(p);
        g[p].push_back(i + 1);
    }

    // construct the LCA
    constexpr int root = 0;
    lowest_common_ancestor lca(root, g);

    // answer to queries
    while (q --) {
        int u, v; scanf("%d%d", &u, &v);
        printf("%d\n", lca(u, v));
    }
    return 0;
}

STACK_PIVOT_MAIN(moin)