Centroid decomposition

centroidDecomp.cpp

// https://codeforces.com/problemset/problem/342/E

#include <bits/stdc++.h>

using namespace std;
typedef int ll;

const ll INF = 1e9;

class Tree {
   public:
    vector<ll> subtree_size;
    vector<set<ll>> adj;

    void calc_subtree_size(ll u, ll p = -1) {
        subtree_size[u] = 1;

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

            calc_subtree_size(v, u);

            subtree_size[u] += subtree_size[v];
        }
    }

    ll centroid(ll u, ll sz, ll p = -1) {
        for (ll v : adj[u]) {
            if (v == p) {
                continue;
            }

            if (subtree_size[v] * 2 > sz) {
                return centroid(v, sz, u);
            }
        }

        return u;
    }

    void add_edge(ll u, ll v) {
        if (u == v) {
            return;
        }

        adj[u].insert(v);
        adj[v].insert(u);
    }

    ll size() {
        return adj.size();
    }

    Tree(ll n) {
        adj.resize(n);
        subtree_size.resize(n, -1);
    }

    Tree() {}
};

class CentroidTree {
    vector<vector<pair<ll, ll>>> adj;
    vector<ll> minDist;  // minimum distance to a red node

    void calcDist(Tree& tree, ll u, ll root, ll p = -1, ll curDist = 1) {
        for (ll v : tree.adj[u]) {
            if (v == p) {
                continue;
            }

            calcDist(tree, v, root, u, curDist + 1);
        }

        adj[u].push_back({root, curDist});
    }

    void decompose(Tree& tree, ll u) {
        tree.calc_subtree_size(u);

        ll x = tree.centroid(u, tree.subtree_size[u]);

        for (ll v : tree.adj[x]) {
            calcDist(tree, v, x, x);
        }

        for (ll v : tree.adj[x]) {
            tree.adj[v].erase(x);
            decompose(tree, v);
        }

        tree.adj[x].clear();
    }

   public:
    void paint(ll u) {
        for (auto [v, w] : adj[u]) {
            minDist[v] = min(minDist[v], w);
        }

        minDist[u] = 0;
    }

    ll query(ll u) {
        ll res = minDist[u];

        for (auto [v, w] : adj[u]) {
            if (!w) {
                continue;
            }

            res = min(res, minDist[v] + w);
        }

        return res;
    }

    CentroidTree(Tree& tree) {
        adj.resize(tree.size());
        decompose(tree, 0);
        minDist.resize(tree.size(), INF);
    }
};

int main() {
    ios::sync_with_stdio(0);
    cout.tie(0);
    cin.tie(0);

    ll n, m;
    cin >> n >> m;

    Tree tree(n);

    for (ll i = 0; i < n - 1; i++) {
        ll u, v;
        cin >> u >> v;

        tree.add_edge(u - 1, v - 1);
    }

    CentroidTree cTree(tree);

    cTree.paint(0);

    while (m--) {
        ll t, v;
        cin >> t >> v;

        if (t == 1) {
            cTree.paint(v - 1);
        } else {
            cout << cTree.query(v - 1) << '\n';
        }
    }
}