【CodeForces-400D】Dima and Bacteria【Disjoint Set + Floyd Warshall Algorithm】

【CodeForces-400D】Dima and Bacteria【Disjoint Set + Floyd Warshall Algorithm】

代码:

#include <iostream>
#include <vector>

using namespace std;
typedef long long LL;
struct Edge {
    int to, w;
};
const int NLIM = 1e5 + 10;
const int SLIM = 500 + 10;
const int INF = 0xfffffff;
vector<Edge> adjOfNode[NLIM];
int parent[NLIM];
int ns[NLIM];       /* Node Set */
int c[SLIM];
int adjOfSet[SLIM][SLIM];

inline int mini(int a, int b) {return a < b ? a : b;}

int Find(int x) {
    if (parent[x] != x) parent[x] = Find(parent[x]);
    return parent[x];
}

void Union(int x, int y) {
    int nx = Find(x);
    int ny = Find(y);
    if (nx != ny) parent[nx] = parent[ny];
}

int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    int n, m, k;
    cin >> n >> m >> k;
    for (int i = 0; i < k; i++) {
        cin >> c[i];
        if (i) c[i] += c[i - 1];
    }
    for (int i = 0; i < k; i++) {
        int crt = i ? c[i - 1] + 1 : 1;
        while (crt <= c[i]) {
            // save each node belong to which set
            ns[crt] = i;
            // init parent array
            parent[crt] = crt;
            crt++;
        }
    }
    // init adjOfSet
    for (int i = 0; i < k; i++) {
        for (int j = 0; j < k; j++) {
            adjOfSet[i][j] = INF;
        }
    }
    // load the adjacent list and matrix
    int u, v;
    Edge tmp;
    for (int i = 0; i < m; i++) {
        cin >> u >> v >> tmp.w;
        tmp.to = v;
        adjOfNode[u].push_back(tmp);
        tmp.to = u;
        adjOfNode[v].push_back(tmp);
        // cout << u << " " << v << " " << ns[u] << " " << ns[v] << " " << tmp.w << endl;
        adjOfSet[ns[u]][ns[v]] = mini(adjOfSet[ns[u]][ns[v]], tmp.w);
        adjOfSet[ns[v]][ns[u]] = mini(adjOfSet[ns[v]][ns[u]], tmp.w);
    }
    // use union-find set to check validity
    bool flag = true;
    for (int i = 0; i < k; i++) {
        int crt = i ? c[i - 1] + 1 : 1;
        if (c[i] == crt) {
            adjOfSet[ns[crt]][ns[crt]] = 0;
        }
        while (crt <= c[i]) {
            for (int j = 0; j < adjOfNode[crt].size(); j++) {
                if (adjOfNode[crt][j].w == 0)
                    Union(crt, adjOfNode[crt][j].to);
            }
            crt++;
        }
        crt = i ? c[i - 1] + 1 : 1;
        int pre = Find(crt++);
        for (; crt <= c[i]; crt++) {
            if (Find(crt) != pre) break;
        }
        if (crt <= c[i]) {
            flag = false;
            break;
        }
    }
    if (flag) {
        // for (int i = 0; i < k; i++) {
        //     for (int j = 0; j < k; j++) {
        //         if (j) cout << ' ';
        //         cout << adjOfSet[i][j];
        //     }
        //     cout << endl;
        // }
        cout << "Yes\n";
        // run the FloydDP
        for (int r = 0; r < k; r++) {
            for (int i = 0; i < k; i++) {
                for (int j = 0; j < k; j++) {
                    adjOfSet[i][j] = mini(adjOfSet[i][j], adjOfSet[i][r] + adjOfSet[r][j]);
                }
            }
        }
        for (int i = 0; i < k; i++) {
            for (int j = 0; j < k; j++) {
                if (j) cout << " ";
                cout << (adjOfSet[i][j] == INF ? -1 : adjOfSet[i][j]);
            }
            cout << "\n";
        }
    } else {
        cout << "No\n";
    }
    return 0;
}