Codeforces Round #385 (Div. 2), problem: (C) Hongcow Builds A Nation Solution in C/C++

#include <stdio.h>
#include <string.h>

int c[1005], p[1005], sz[1005], gov[1005];

int find_parent(int n) {
    if (p[n] == n)
        return n;
    return p[n] = find_parent(p[n]);
}

int get_full_size(int n) {
    return n * (n-1) / 2;
}

int main(int argc, char** argv) {
    int i, n, m, k, u, v, max_size, total_free, used, total_edge;

    scanf("%d%d%d", &n, &m, &k);
    for (i = 0; i < k; i++)
        scanf("%d", &c[i]);

    for (i = 1; i <= n; i++) {
        p[i] = i;
        sz[i] = 1;
    }

    for (i = 0; i < m; i++) {
        scanf("%d%d", &u, &v);
        if (find_parent(u) != find_parent(v)) {
            sz[p[v]] += sz[p[u]];
            sz[p[u]] = 0;
            p[p[u]] = p[v];
        }
    }

    memset(gov, 0, sizeof gov);
    max_size = 0;
    for (i = 0; i < k; i++) {
        c[i] = find_parent(c[i]);
        gov[c[i]] = 1;
        if (sz[c[i]] > max_size)
            max_size = sz[c[i]];
    }

    total_free = 0;
    for (i = 1; i <= n; i++)
        if (!gov[i])
            total_free += sz[i];

    used = 0;
    total_edge = 0;
    for (i = 0; i < k; i++) {
        if (sz[c[i]] == max_size && !used) {
            total_edge += get_full_size(max_size+total_free);
            used = 1;
        }
        else
            total_edge += get_full_size(sz[c[i]]);
    }

    printf("%d\n", total_edge - m);

    return 0;
}

Leave a Reply

Your email address will not be published. Required fields are marked *