Heavy light decomposition - YessineJallouli/Competitive-Programming GitHub Wiki

#include <bits/stdc++.h>
#define ll long long
#define SaveTime ios_base::sync_with_stdio(false), cin.tie(0);
using namespace std;

const int N = 1e5+7;
int sz[N], heavy[N], depth[N], top[N], pos[N], par[N];
vector<vector<int>> tree(N);
vector<int> segment_ids;
int n;

void dfs(int node, int pr = -1) {
    sz[node] = 1;
    int mx = 0;
    for (int ch : tree[node]) {
        if (ch == pr)
            continue;
        par[ch] = node;
        depth[ch] = depth[node] + 1;
        dfs(ch, node);
        sz[node]+= sz[ch];
        if (mx < sz[ch]) {
            mx = sz[ch];
            heavy[node] = ch;
        }
    }
}

void decompose(int node, int head, int pr = -1) {
    segment_ids.push_back(node);
    pos[node] = segment_ids.size()-1;
    top[node] = head;
    if (heavy[node]) {
        decompose(heavy[node], head, node);
    }
    for (int ch : tree[node]) {
        if (ch == pr || ch == heavy[node])
            continue;
        decompose(ch, ch, node);
    }
}

int main() {
    //SaveTime;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        int nb; cin >> nb;
        for (int j = 0; j < nb; j++) {
            int v; cin >> v;
            tree[i].push_back(v);
            tree[v].push_back(i);
        }
    }
    depth[1] = 0;
    par[1] = 1;
    dfs(1);
    decompose(1, 1);
    int q; cin >> q;
    while (q--) {
        int a,b; cin >> a >> b;
        for (; top[a] != top[b]; b = par[top[b]]) {
            if (depth[top[a]] > depth[top[b]])
                swap(a,b);
            // query : pos[top[b]] .. pos[b]
        }
        if (depth[a] > depth[b])
            swap(a,b);
        // query : pos[a] .. pos[b]
        cout << a << '\n';
    }
}

#include <bits/stdc++.h>
#define SaveTime ios_base::sync_with_stdio(false), cin.tie(0);
using namespace std;

const int N = 1e5+7;
int values[N];
int sz[N], heavy[N], depth[N], top[N], pos[N], par[N];
vector<vector<pair<int,int>>> tree(N);
vector<int> segment_ids;
pair<int,int> edgeOrder[N];

int n;
int segTree[4*N];

int mrg(int x, int y) {
    // segTree mrg
}

void build(int id = 0, int ns = 0, int ne = n-1) {
    if (ns == ne) {
        segTree[id] = values[segment_ids[ns]];
        return;
    }
    int l = 2*id+1;
    int r = 2*id+2;
    int md = (ns+ne)/2;
    build(l, ns, md);
    build(r, md+1, ne);
    segTree[id] = mrg(segTree[l], segTree[r]);
}

int get(int qs, int qe, int id = 0, int ns = 0, int ne = n-1) {
    // segTree get
}

void upd(int pos, int val, int id = 0, int ns = 0, int ne = n-1) {
    // segTree upd
}


void dfs(int node, int val = 0, int pr = -1) {
    sz[node] = 1;
    values[node] = val;
    int mx = 0;
    for (auto [ch,w] : tree[node]) {
        if (ch == pr)
            continue;
        par[ch] = node;
        depth[ch] = depth[node] + 1;
        dfs(ch, w, node);
        sz[node]+= sz[ch];
        if (mx < sz[ch]) {
            mx = sz[ch];
            heavy[node] = ch;
        }
    }
}

void decompose(int node, int head, int pr = -1) {

    segment_ids.push_back(node);
    pos[node] = (int) segment_ids.size()-1;
    top[node] = head;

    if (heavy[node])
        decompose(heavy[node], head, node);

    for (auto [ch,w] : tree[node]) {
        if (ch == pr || ch == heavy[node])
            continue;
        decompose(ch, ch, node);
    }
}

int query_HLD(int a, int b) {
    int ans = 0;
    for (; top[a] != top[b]; b = par[top[b]]) {
        if (depth[top[a]] > depth[top[b]])
            swap(a, b);
        ans = mrg(ans, get(pos[top[b]], pos[b]));
    }
    if (depth[a] > depth[b])
        swap(a, b);
    ans = mrg(ans, get(pos[a]+1, pos[b]));
    return ans;
}

void upd_HLD(int idEdge, int val) {
    auto [node1, node2] = edgeOrder[idEdge];
    if (depth[node1] < depth[node2])
        swap(node1, node2);
    upd(pos[node1], val);
}

void init() {
    segment_ids.clear();
    for (int i = 0; i <= n; i++) {
        heavy[i] = 0;
        tree[i].clear();
    }
}

int main() {
    SaveTime
    int T; cin >> T;
    while (T--) {
        cin >> n;
        init();

        for (int i = 0; i < n-1; i++) {
            int x,y,w; cin >> x >> y >> w;
            tree[x].emplace_back(y,w);
            tree[y].emplace_back(x,w);
            edgeOrder[i] = {x,y};
        }

        depth[1] = 0;
        par[1] = 1;
        dfs(1);
        decompose(1, 1);
        build();

        string s; cin >> s;
        while (s !=  "DONE") {
            int a, b; cin >> a >> b;

            if (s == "CHANGE")
                upd_HLD(a-1, b);

            else if (s == "QUERY")
                cout << query_HLD(a,b) << '\n';

            cin >> s;
        }
    }
}