atcoder/hld.hpp

• View this file on GitHub
• Last update: 2021-10-21 18:06:39+09:00
• Include: #include "atcoder/hld.hpp"
• Link: https://ei1333.github.io/luzhiled/snippets/tree/heavy-light-decomposition.html

Verified with

• test/hld.test.cpp
• test/persistentsegtree.test.cpp

Code

#ifndef ATCODER_HLD_HPP
#define ATCODER_HLD_HPP 1

#include <vector>

namespace atcoder {

// Reference: https://ei1333.github.io/luzhiled/snippets/tree/heavy-light-decomposition.html

template< typename G >
struct HeavyLightDecomposition {
    G &g;
    std::vector< int > sz, in, out, head, rev, par;

    explicit HeavyLightDecomposition(G &g_) :
            g(g_), sz(g.size()), in(g.size()), out(g.size()), head(g.size()), rev(g.size()), par(g.size()) {
        build();
    }

    int la(int v, int k) {
        while(1) {
            int u = head[v];
            if(in[v] - k >= in[u]) return rev[in[v] - k];
            k -= in[v] - in[u] + 1;
            v = par[u];
        }
    }

    int lca(int u, int v) {
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) return u;
        }
    }

    template< typename T, typename Q, typename F >
    T prod(int u, int v, const T &id, const Q &q, const F &f, bool edge = false) {
        // Require: f(a, b) == f(b, a)
        T s = id;
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) break;
            s = f(s, q(in[head[v]], in[v] + 1));
        }
        return f(s, q(in[u] + edge, in[v] + 1));
    }

    template< typename Q >
    void apply(int u, int v, const Q &q, bool edge = false) {
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) break;
            q(in[head[v]], in[v] + 1);
        }
        q(in[u] + edge, in[v] + 1);
    }

private:
    void dfs_sz(int idx, int p) {
        par[idx] = p;
        sz[idx] = 1;
        if(g[idx].size() && g[idx][0] == p) std::swap(g[idx][0], g[idx].back());
        for(auto &to : g[idx]) {
            if(to == p) continue;
            dfs_sz(to, idx);
            sz[idx] += sz[to];
            if(sz[g[idx][0]] < sz[to]) std::swap(g[idx][0], to);
        }
    }

    void dfs_hld(int idx, int p, int &times) {
        in[idx] = times++;
        rev[in[idx]] = idx;
        for(auto &to : g[idx]) {
            if(to == p) continue;
            head[to] = (g[idx][0] == to ? head[idx] : to);
            dfs_hld(to, idx, times);
        }
        out[idx] = times;
    }

    void build() {
        dfs_sz(0, -1);
        int t = 0;
        dfs_hld(0, -1, t);
    }
};

}  // namespace atcoder

#endif  // ATCODER_HLD_HPP

#line 1 "atcoder/hld.hpp"



#include <vector>

namespace atcoder {

// Reference: https://ei1333.github.io/luzhiled/snippets/tree/heavy-light-decomposition.html

template< typename G >
struct HeavyLightDecomposition {
    G &g;
    std::vector< int > sz, in, out, head, rev, par;

    explicit HeavyLightDecomposition(G &g_) :
            g(g_), sz(g.size()), in(g.size()), out(g.size()), head(g.size()), rev(g.size()), par(g.size()) {
        build();
    }

    int la(int v, int k) {
        while(1) {
            int u = head[v];
            if(in[v] - k >= in[u]) return rev[in[v] - k];
            k -= in[v] - in[u] + 1;
            v = par[u];
        }
    }

    int lca(int u, int v) {
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) return u;
        }
    }

    template< typename T, typename Q, typename F >
    T prod(int u, int v, const T &id, const Q &q, const F &f, bool edge = false) {
        // Require: f(a, b) == f(b, a)
        T s = id;
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) break;
            s = f(s, q(in[head[v]], in[v] + 1));
        }
        return f(s, q(in[u] + edge, in[v] + 1));
    }

    template< typename Q >
    void apply(int u, int v, const Q &q, bool edge = false) {
        for(;; v = par[head[v]]) {
            if(in[u] > in[v]) std::swap(u, v);
            if(head[u] == head[v]) break;
            q(in[head[v]], in[v] + 1);
        }
        q(in[u] + edge, in[v] + 1);
    }

private:
    void dfs_sz(int idx, int p) {
        par[idx] = p;
        sz[idx] = 1;
        if(g[idx].size() && g[idx][0] == p) std::swap(g[idx][0], g[idx].back());
        for(auto &to : g[idx]) {
            if(to == p) continue;
            dfs_sz(to, idx);
            sz[idx] += sz[to];
            if(sz[g[idx][0]] < sz[to]) std::swap(g[idx][0], to);
        }
    }

    void dfs_hld(int idx, int p, int &times) {
        in[idx] = times++;
        rev[in[idx]] = idx;
        for(auto &to : g[idx]) {
            if(to == p) continue;
            head[to] = (g[idx][0] == to ? head[idx] : to);
            dfs_hld(to, idx, times);
        }
        out[idx] = times;
    }

    void build() {
        dfs_sz(0, -1);
        int t = 0;
        dfs_hld(0, -1, t);
    }
};

}  // namespace atcoder

Back to top page