Heavy Light Decomposition

hld.cpp

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;

class SegmentTree {
   private:
    struct Node {
        ll value;
        pair<ll, ll> range;
        Node* left = nullptr;
        Node* right = nullptr;

        ~Node() {
            delete left;
            delete right;
        }
    };

    ll size;
    Node* root = new Node;

    void update_main(Node* cur, ll i, ll val) {
        auto [l, r] = cur->range;

        if (!(l <= i && i <= r)) {
            return;
        }

        if (l == r) {
            cur->value = val;
            return;
        }

        if (cur->left == nullptr) {
            cur->left = new Node;
            cur->left->range = {l, (l + r) / 2};
        }

        if (cur->right == nullptr) {
            cur->right = new Node;
            cur->right->range = {(l + r) / 2 + 1, r};
        }

        update_main(cur->left, i, val);
        update_main(cur->right, i, val);

        cur->value = max(cur->left->value, cur->right->value);
    }

    ll query_main(Node* cur, ll start, ll end) {
        if (cur == nullptr) {
            return 0;
        }

        auto [l, r] = cur->range;

        if (r < start || end < l) {
            return 0;
        }

        if (start <= l && r <= end) {
            return cur->value;
        }

        return max(query_main(cur->left, start, end), query_main(cur->right, start, end));
    }

   public:
    SegmentTree() {}

    SegmentTree(ll n) {
        size = n - 1;
        root->range = {0, size};
    }

    ~SegmentTree() {
        delete root;
    }

    void update(ll i, ll val) {
        update_main(root, i, val);
    }

    ll query(ll l, ll r) {
        return query_main(root, l, r);
    }
};

vector<ll> idx;

struct Component {
    SegmentTree* tree;
    ll root;
    ll size = 0;

    Component() {}

    ~Component() {
        delete tree;
    }

    void init() {
        tree = new SegmentTree(size);
    }

    void push_back(ll val) {
        if (size == 0) {
            root = val;
        }

        idx[val] = size++;
    }

    void set_value(ll item, ll val) {
        tree->update(idx[item], val);
    }

    ll query(ll a, ll b) {
        if (idx[a] > idx[b]) {
            swap(a, b);
        }

        return tree->query(idx[a], idx[b]);
    }

    ll queryRest(ll a) {
        return tree->query(idx[a], size);
    }
};

class HLDTree {
   private:
    vector<Component*> comp;
    vector<ll> parent, subtreeSize, ids, depth;

    ll calc_subtree(ll u, ll p, ll d) {
        parent[u] = p;
        depth[u] = d;
        ll res = 1;

        for (ll i : adj[u]) {
            if (i == p) {
                continue;
            }

            res += calc_subtree(i, u, d + 1);
        }

        return subtreeSize[u] = res;
    }

    void sort_by_size() {
        for (auto& i : adj) {
            vector<pair<ll, ll>> v;

            for (ll j : i) {
                v.push_back({subtreeSize[j], j});
            }

            sort(v.begin(), v.end(), greater{});

            for (ll j = 0; j < i.size(); j++) {
                i[j] = v[j].second;
            }
        }
    }

    void calc_HLD(ll u, ll id) {
        bool isFirst = true;
        ids[u] = id;
        comp[id]->push_back(u);

        for (ll i : adj[u]) {
            if (i == parent[u]) {
                continue;
            }

            if (isFirst) {
                isFirst = false;
                calc_HLD(i, id);
            } else {
                comp.push_back(new Component());
                calc_HLD(i, comp.size() - 1);
            }
        }
    }

    void init_comp() {
        for (auto i : comp) {
            i->init();
        }
    }

   public:
    vector<vector<ll>> adj;
    ll root;

    HLDTree(ll n) {
        adj = vector<vector<ll>>(n);
        parent = subtreeSize = ids = idx = depth = vector<ll>(n);
        comp = vector<Component*>(1, new Component());
    }

    ~HLDTree() {
        for (auto i : comp) {
            delete i;
        }
    }

    void add_edge(ll a, ll b) {
        adj[a].push_back(b);
    }

    void set_root(ll x) {
        root = x;

        calc_subtree(root, 0, 0);
        sort_by_size();
        calc_HLD(root, 0);
        init_comp();
    }

    void set_value(ll node, ll val) {
        comp[ids[node]]->set_value(node, val);
    }

    ll lca(ll a, ll b) {
        while (ids[a] != ids[b]) {
            if (depth[comp[ids[a]]->root] > depth[comp[ids[b]]->root]) {
                swap(a, b);
            }

            b = parent[comp[ids[b]]->root];
        }

        return (depth[a] < depth[b] ? a : b);
    }

    ll query_path(ll a, ll b) {
        ll res = INT_MIN;

        while (ids[a] != ids[b]) {
            if (depth[comp[ids[a]]->root] > depth[comp[ids[b]]->root]) {
                swap(a, b);
            }

            auto& x = comp[ids[b]];
            res = max(res, x->query(b, x->root));
            b = parent[x->root];
        }

        res = max(res, comp[ids[a]]->query(a, b));

        return res;
    }

    ll query_subtree(ll a) {
        ll ans = comp[ids[a]]->queryRest(a);

        for (auto& i : comp) {
            if (ids[a] == ids[i->root]) {
                continue;
            }

            if (lca(a, i->root) == a) {
                ans = max(ans, i->queryRest(i->root));
            }
        }

        return ans;
    }
};

int main() {

}