test/yosupo/lca_doubling.test.cpp

• View this file on GitHub
• Last update: 2023-07-03 22:12:45+09:00
• Problem: https://judge.yosupo.jp/problem/lca

Depends on

• tree/doubling_tree.hpp

Code

#include <iostream>
#include <optional>
#include <vector>

#include "../../tree/doubling_tree.hpp"

#define PROBLEM "https://judge.yosupo.jp/problem/lca"

#define _overload(_1, _2, _3, _4, name, ...) name
#define _rep1(Itr, N) _rep3(Itr, 0, N, 1)
#define _rep2(Itr, A, B) _rep3(Itr, A, B, 1)
#define _rep3(Itr, A, B, S) for (i64 Itr = A; Itr < B; Itr += S)
#define repeat(...) _overload(__VA_ARGS__, _rep3, _rep2, _rep1)(__VA_ARGS__)
#define rep(...) _overload(__VA_ARGS__, _rep3, _rep2, _rep1)(__VA_ARGS__)
#define ALL(X) begin(X), end(X)

using i64 = long long;
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    i64 n, q;
    cin >> n >> q;
    vector<optional<i64>> p(n, nullopt);
    rep(i, 1, n) {
        i64 v;
        cin >> v;
        p[i] = v;
    }

    DoublingTree lca(ALL(p));

    rep(i, q) {
        i64 u, v;
        cin >> u >> v;
        cout << lca.lca(u, v) << '\n';
    }

    return 0;
};

#line 1 "test/yosupo/lca_doubling.test.cpp"
#include <iostream>
#include <optional>
#include <vector>

#line 1 "tree/doubling_tree.hpp"


#include <iterator>
#line 6 "tree/doubling_tree.hpp"

// 0-indexed
// climb(u, d): climb d steps towards root
// fold(u, v):
struct DoublingTree {
    std::vector<std::vector<std::optional<int>>> parent;
    std::vector<int> depth;
    int logn;

    template <class InputItr>
    DoublingTree(InputItr first, InputItr last) {
        int n = std::distance(first, last);
        std::vector<std::optional<int>> p(n, std::nullopt);
        int i = 0;
        for (auto itr = first; itr != last; itr++, i++) {
            if (itr->has_value()) p[i] = (int)itr->value();
        }
        build(p);
    };

    void build(const std::vector<std::optional<int>> &p) {
        int n = p.size();
        logn = 1;
        while ((1 << logn) < n) logn++;

        parent.assign(logn, std::vector<std::optional<int>>(n, std::nullopt));
        for (int i = 0; i < n; i++) parent[0][i] = p[i];
        std::vector<std::vector<int>> tree(n);
        std::vector<int> root;
        for (int i = 0; i < n; i++) {
            if (parent[0][i].has_value())
                tree[parent[0][i].value()].push_back(i);
            else
                root.push_back(i);
        }

        depth.assign(n, -1);
        auto calc_depth = [&](int u, int d, auto &&f) -> void {
            depth[u] = d;
            for (auto v : tree[u]) {
                f(v, d + 1, f);
            }
            return;
        };

        for (int u : root) {
            calc_depth(u, 0, calc_depth);
        }

        for (int k = 1; k < logn; k++) {
            for (int u = 0; u < n; u++) {
                if (parent[k - 1][u].has_value())
                    parent[k][u] = parent[k - 1][parent[k - 1][u].value()];
            }
        }
    };

    std::optional<int> climb(int u, int d) {
        if (d > depth[u]) return std::nullopt;
        int cnt = 0;
        while (d > 0) {
            if (d & 1) u = parent[cnt][u].value();
            d >>= 1;
            cnt++;
        }
        return u;
    };

    // LowestCommonAncestor
    int lca(int u, int v) {
        if (depth[u] > depth[v]) std::swap(u, v);
        v = climb(v, depth[v] - depth[u]).value();

        if (u == v) return u;
        for (int k = logn - 1; k >= 0; k--) {
            if (parent[k][u] != parent[k][v]) {
                u = parent[k][u].value();
                v = parent[k][v].value();
            }
        }
        return parent[0][u].value();
    };

    inline int distance(int u, int v) {
        return depth[u] + depth[v] - depth[lca(u, v)] * 2;
    };
};

struct DoublingTreeBuilder {
    std::vector<std::vector<int>> g;
    DoublingTreeBuilder(int n) : g(n){};
    void add_edge(int a, int b) {
        g[a].push_back(b);
        g[b].push_back(a);
    };

    DoublingTree build(const std::vector<int> &root = {0}) {
        std::vector<std::optional<int>> parent(g.size(), std::nullopt);
        auto dfs = [&](int u, int p, auto &&f) -> void {
            for (auto v : g[u]) {
                if (v == p) continue;
                parent[v] = u;
                f(v, u, f);
            }
            return;
        };
        for (auto v : root) dfs(v, -1, dfs);
        return DoublingTree(parent.begin(), parent.end());
    };
};


#line 6 "test/yosupo/lca_doubling.test.cpp"

#define PROBLEM "https://judge.yosupo.jp/problem/lca"

#define _overload(_1, _2, _3, _4, name, ...) name
#define _rep1(Itr, N) _rep3(Itr, 0, N, 1)
#define _rep2(Itr, A, B) _rep3(Itr, A, B, 1)
#define _rep3(Itr, A, B, S) for (i64 Itr = A; Itr < B; Itr += S)
#define repeat(...) _overload(__VA_ARGS__, _rep3, _rep2, _rep1)(__VA_ARGS__)
#define rep(...) _overload(__VA_ARGS__, _rep3, _rep2, _rep1)(__VA_ARGS__)
#define ALL(X) begin(X), end(X)

using i64 = long long;
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    i64 n, q;
    cin >> n >> q;
    vector<optional<i64>> p(n, nullopt);
    rep(i, 1, n) {
        i64 v;
        cin >> v;
        p[i] = v;
    }

    DoublingTree lca(ALL(p));

    rep(i, q) {
        i64 u, v;
        cin >> u >> v;
        cout << lca.lca(u, v) << '\n';
    }

    return 0;
};

Back to top page