centroid_decomposition.cpp

/*16
1 4
2 4
3 4
4 5
5 6
6 7
7 8
7 9
6 10
10 11
11 12
12 14
11 13
13 15
13 16*/

#include <bits/stdc++.h>

#define N 100001

std::vector<std::vector<int>> adj;
std::vector<std::vector<int>> centroid_tree;

int subsizes[N];
bool removed[N];

/* This function computes the sizes of subtrees starting from node */
int DFS_sizes(int node, int par = -1)
{
    subsizes[node] = 1;
    
    for (auto& elem: adj[node])
        if (!removed[elem] && elem != par)
            subsizes[node] += DFS_sizes(elem, node);
        
    return subsizes[node];
}

/*After sizes of subtrees are calculated, we can find the centroid */
int get_centroid(int node, int total_sz, int par = -1)
{
    for (auto& elem : adj[node])
        if (!removed[elem] && elem != par)
            if (subsizes[elem] > total_sz / 2)
                return get_centroid(elem, total_sz, node);
            
    return node;
}

/* Performs Centroid decomposition of a tree */
int cent_decompose(int node)
{
    int total_sz = DFS_sizes(node, -1);
    
    if (total_sz == 1) // Leaf node on centroid tree
    {
        removed[node] = true;
        return node;
    }
    else
    {
        int centroid = get_centroid(node, total_sz, -1);
        removed[centroid] = true;
        
        // Populate centroid tree
        for (auto& elem : adj[centroid])
            if (!removed[elem])
                centroid_tree[centroid].push_back(cent_decompose(elem));
        
        return centroid;
    }
}

int main()
{
    int n;
    std::cin >> n;
    
    adj.resize(n);
    centroid_tree.resize(n);
    std::fill(removed, removed + n, false);
    
    // Input
    for (int i = 1; i < n; ++i)
    {
        int l, r;
        std::cin >> l >> r;
        l--; r--;
        
        adj[l].push_back(r);
        adj[r].push_back(l);
    }
    
    
    int root = cent_decompose(0);
    
    std::cout << "Root: " << root << std::endl;
    std::cout << "Centroid Tree Adjacency List: " << std::endl;
    for (int i = 0; i < centroid_tree.size(); ++i)
    {
        std::cout << i << ": ";
        for (auto& elem: centroid_tree[i])
            std::cout << elem << " ";
        
        std::cout << std::endl;
    }
    
    return 0;
}