[CF400D] Dima and Bacteria - 最短路,并查集

[CF400D] Dima and Bacteria - 最短路,并查集

Description

有n个点，每个点都有自己的type，给出m条边，首先要求判断type相同的点之间能否以0为代价互相连通，如果可以，求任意两种type连通的最小代价

Solution

把所有 0 的边都用并查集玩一下，看每个 type 的点是否处在同一个连通块中

因为 \(k \le 500\)，我们可以把每个 type 缩成一个点，然后重新建图，跑 floyd

#include <bits/stdc++.h>
using namespace std;

#define int long long

const int N = 100005;
int n, m, k;
int c[N];
int x[N];
int bel[N];
int fa[N];
int d[505][505];

int find(int p)
{
    return p == fa[p] ? p : fa[p] = find(fa[p]);
}

void merge(int p, int q)
{
    p = find(p);
    q = find(q);
    if (p - q)
        fa[p] = q;
}

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n >> m >> k;
    for (int i = 1; i <= k; i++)
        cin >> c[i];
    int pos = 1;
    for (int i = 1; i <= k; i++)
    {
        for (int j = pos; j < pos + c[i]; j++)
            bel[j] = i;
        pos += c[i];
    }

    vector<tuple<int, int, int>> edges;
    for (int i = 1; i <= m; i++)
    {
        int t1, t2, t3;
        cin >> t1 >> t2 >> t3;
        edges.push_back({t1, t2, t3});
    }

    for (int i = 1; i <= n; i++)
        fa[i] = i;

    for (auto [u, v, w] : edges)
    {
        if (w == 0)
            merge(u, v);
    }

    pos = 1;
    for (int i = 1; i <= k; i++)
    {
        for (int j = pos; j < pos + c[i]; j++)
            if (find(j) != find(pos))
            {
                cout << "No" << endl;
                return 0;
            }
        pos += c[i];
    }

    cout << "Yes" << endl;

    int ind = k;

    for (auto &[u, v, w] : edges)
    {
        u = bel[u];
        v = bel[v];
    }

    memset(d, 0x3f, sizeof d);
    for (auto [u, v, w] : edges)
    {
        d[u][v] = min(d[u][v], w);
        d[v][u] = min(d[v][u], w);
    }

    for (int i = 1; i <= ind; i++)
    {
        d[i][i] = 0;
    }

    for (int k = 1; k <= ind; k++)
    {
        for (int i = 1; i <= ind; i++)
        {
            for (int j = 1; j <= ind; j++)
            {
                d[i][j] = min(d[i][j], d[i][k] + d[k][j]);
            }
        }
    }

    for (int i = 1; i <= ind; i++)
    {
        for (int j = 1; j <= ind; j++)
        {
            if (d[i][j] < 1e12)
                cout << d[i][j] << " ";
            else
                cout << -1 << " ";
        }
        cout << endl;
    }
}