Heavy-light decomposition

HLD.cpp

#include <bits/stdc++.h>

using namespace std;

template<typename _Ty, typename _Fn = plus<_Ty>>
struct segment_tree {
    _Ty default_value = _Ty();
    int n = 0;
    vector<_Ty> tree;

    void _init(int x, int s, int e, const vector<_Ty> &a) {
        if (s == e) {
            tree[x] = a[s];
            return;
        }
        _init(x * 2, s, (s + e) / 2, a);
        _init(x * 2 + 1, (s + e) / 2 + 1, e, a);
        tree[x] = _Fn()(tree[x * 2], tree[x * 2 + 1]);
    }

    segment_tree() : n(0) {}

    explicit segment_tree(const vector<_Ty> &a, _Ty default_value = _Ty()) : default_value(default_value) {
        n = a.size();
        int ts = 1;
        while (ts < a.size()) ts <<= 1;
        ts <<= 1;
        tree = vector<_Ty>(ts, default_value);
        _init(1, 0, n - 1, a);
    }

    _Ty _query(int x, int s, int e, int l, int r) {
        if (r < s || e < l) return default_value;
        if (l <= s && e <= r) return tree[x];
        return _Fn()(
                _query(x * 2, s, (s + e) / 2, l, r),
                _query(x * 2 + 1, (s + e) / 2 + 1, e, l, r)
        );
    }

    _Ty query(int l, int r) {
        return _query(1, 0, n - 1, l, r);
    }

    void _change(int x, int s, int e, int i, _Ty val) {
        if (e < i || i < s) return;
        if (s == e && e == i) {
            tree[x] = val;
            return;
        }
        _change(x * 2, s, (s + e) / 2, i, val);
        _change(x * 2 + 1, (s + e) / 2 + 1, e, i, val);
        tree[x] = _Fn()(tree[x * 2], tree[x * 2 + 1]);
    }

    void change(int i, _Ty val) {
        _change(1, 0, n - 1, i, val);
    }
};

template<typename _Ty, typename _Fn = plus<_Ty>>
struct hld {
    vector<int> parent, depth, heavy, head, pos;
    vector<_Ty> values, segmented_values;
    segment_tree<_Ty, _Fn> tree;
    int cur_pos = 0;
    _Ty default_value = _Ty();

    int dfs(int v, const vector<vector<int>> &adj) {
        int size = 1;
        int max_c_size = 0;
        for (const auto &c : adj[v]) {
            if (c == parent[v]) continue;
            parent[c] = v, depth[c] = depth[v] + 1;
            int c_size = dfs(c, adj);
            size += c_size;
            if (c_size > max_c_size)
                max_c_size = c_size, heavy[v] = c;
        }
        return size;
    }

    void decompose(int v, int h, const vector<vector<int>> &adj) {
        head[v] = h, pos[v] = cur_pos++;
        segmented_values[pos[v]] = values[v];
        if (heavy[v] != -1)
            decompose(heavy[v], h, adj);
        for (const auto &c : adj[v]) {
            if (c != parent[v] && c != heavy[v]) {
                decompose(c, c, adj);
            }
        }
    }

    hld() = default;

    explicit hld(const vector<vector<int>> &adj, vector<_Ty> values, _Ty default_value = _Ty())
            : values(move(values)), default_value(default_value) {
        int n = adj.size();
        parent = vector<int>(n);
        depth = vector<int>(n);
        heavy = vector<int>(n, -1);
        head = vector<int>(n);
        pos = vector<int>(n);
        segmented_values = vector<_Ty>(n);
        cur_pos = 0;

        dfs(0, adj);
        decompose(0, 0, adj);
        tree = segment_tree<_Ty, _Fn>(segmented_values, default_value);
    }

    _Ty query(int a, int b) {
        int res = default_value;
        for (; head[a] != head[b]; b = parent[head[b]]) {
            if (depth[head[a]] > depth[head[b]])
                swap(a, b);
            _Ty cur_heavy_path_max = tree.query(pos[head[b]], pos[b]);
            res = _Fn()(res, cur_heavy_path_max);
        }
        if (depth[a] > depth[b])
            swap(a, b);
        if (pos[a] != pos[b]) {
            _Ty last_heavy_path_max = tree.query(pos[a] + 1, pos[b]);
            res = _Fn()(res, last_heavy_path_max);
        }
        return _Fn()(values[a], res);
    }

    void change(int i, _Ty v) {
        values[i] = v;
        tree.change(pos[i], v);
    }
};