bzoj1086: [SCOI2005]王室联邦

#include <iostream>
#include <cstdlib>
#include <cstdio>
#include <stack>

using namespace std;
const int maxn = 2010, maxm = 8010;

int n, B;

class Graph
{
public:
    struct edge { int u, v; edge *next; };
    edge epool[maxm], *edges[maxn];
    int cnt, root, par[maxn], size[maxn];
    void addedge(int u, int v)
    {
        edge *p = &epool[++cnt], *q = &epool[++cnt];
        p->u = u; p->v = v; p->next = edges[u]; edges[u] = p;
        q->u = v; q->v = u; q->next = edges[v]; edges[v] = q;
        return ;
    }
    // These are the specific operations which would be applied onto the
    // province division procedure.
    stack<int> cities; // How many cities are pending province division
    int city_belong[maxn], // The province ID of the city
        capital[maxn], // The capital of the province, in city ID
        provinces; // The count of provinces, ID starting from 1
    void dfs(int p)
    {
        size[p] = 0;
        cities.push(p);
        for (edge *ep = edges[p]; ep; ep = ep->next)
            if (ep->v != par[p]) {
                par[ep->v] = p;
                dfs(ep->v);
                size[p] += size[ep->v];
                // If the size is sufficient (n > B), we may consume it and not
                // count it as a subtree component.
                if (size[p] >= B) {
                    size[p] = 0;
                    capital[++provinces] = p;
                    int cit;
                    // Pop all children of p that were not popped out before,
                    // these forms the cities of the new province.
                    while ((cit = cities.top()) != p) { // Take all cities
                        cities.pop();
                        city_belong[cit] = provinces; // Set province ID
                    }
                }
            }
        // You must count oneself!
        size[p]++;
        return ;
    }
    // After secure investigation of the function dfs(), we may discover that
    // most capitals do not have belongings, where they are not marked available
    // as a component. We now format them to belong in at some province.
    void format_capitals(int p, int province)
    {
        if (city_belong[p])
            province = city_belong[p];
        else
            city_belong[p] = province;
        // Iterate children
        for (edge *ep = edges[p]; ep; ep = ep->next)
            if (par[ep->v] == p)
                format_capitals(ep->v, province);
        return ;
    }
} graph;

int main(int argc, char** argv)
{
    scanf("%d%d", &n, &B);
    // The form is instinctively and factly a tree, according to description.
    // The capitol of the provinces are either the root of the subtree or the
    // parent of the root of the subtree. This makes division of the tree
    // arbitrarily easy. Therefore we may always find a way to satisfy the
    // range [B, 3B] easily, given n is sufficient.
    if (n < B) {
        printf("0\n");
        return 0;
    }
    // Afterwards we read in the data, formally.
    for (int i = 1, a, b; i <= n - 1; i++) {
        scanf("%d%d", &a, &b);
        graph.addedge(a, b);
    }
    // Deep First Search to generate the tree and calculate subtree size.
    graph.dfs(1);
    // Reformat capitals to look right
    graph.format_capitals(1, graph.provinces);
    // Output answer.
    printf("%d\n", graph.provinces);
    for (int i = 1; i <= n; i++)
        printf("%d ", graph.city_belong[i]);
    printf("\n");
    for (int i = 1; i <= graph.provinces; i++)
        printf("%d ", graph.capital[i]);
    printf("\n");
    return 0;
}