XV Open Cup named after E.V. Pankratiev. GP of Central Europe (AMPPZ-2014)--B.Petrol

多源最短路+并查集

#include <bits/stdc++.h>
using namespace std;
#define rep(i, j, k) for (int i = int(j); i <= int(k); ++ i)
typedef pair<int, int> P;
const int N = 2e5 + 7;
const int inf = 2e9 + 7;
vector<P> g[N];
int c[N], d[N], closest[N];
struct Query {
    int x, y, b, ans, id;
    void read() {
        scanf("%d%d%d", &x, &y, &b);
    }
}b[N];
struct triple {
    int u, v, w;
};
int fa[N];
int findx(int x) {
    if (x == fa[x]) return x; 
    return fa[x] = findx(fa[x]);
}
int main() {
    int n, m, s;
    scanf("%d%d%d", &n, &s, &m);
    rep(i, 1, s) scanf("%d", &c[i]);
    vector<triple> e;
    rep(i, 1, m) {
        int u, v, d;
        scanf("%d%d%d", &u, &v, &d);
        g[u].push_back(P(v, d));
        g[v].push_back(P(u, d));
        e.push_back((triple){u, v, d});
    }
    auto dijkstra = [&]() {
        priority_queue<P, vector<P>, greater<P> > Q;
        rep(i, 1, n) d[i] = inf;
        rep(i, 1, s) {
            d[c[i]] = 0; closest[c[i]] = c[i];
            Q.push(P(0, c[i]));
        }
        while (!Q.empty()) {
            P r = Q.top(); Q.pop();
            int u = r.second;
            if (d[u] < r.first) continue;
            for (P p: g[u]) {
                int v = p.first;
                if (d[v] > d[u] + p.second) {
                    d[v] = d[u] + p.second;
                    closest[v] = closest[u];
                    Q.push(P(d[v], v));
                }
            }
        }
    };
    dijkstra();
    int q;
    scanf("%d", &q);
    rep(i, 1, q) b[i].read(), b[i].id = i;
    auto calc = [&](triple &x) -> int{
        return d[x.u] + d[x.v] + x.w;
    };
    sort(b + 1, b + 1 + q, [&](Query &a, Query &b) {
        return a.b < b.b;
    });
    sort(begin(e), end(e), [&](triple &a, triple &b) {
        return calc(a) < calc(b);
    });
    rep(i, 1, n) fa[i] = i;
    auto same = [&](int x, int y) -> bool {
        x = findx(x); y = findx(y);
        return x == y;
    };
    auto unite = [&](int x, int y) {
        x = findx(x); y = findx(y);
        fa[x] = y;
    };
    int l = 0;
    rep(i, 1, q) {
        while (l < m && calc(e[l]) <= b[i].b) {
            unite(e[l].u, e[l].v);
            l ++;
        }
        if (same(b[i].x, b[i].y)) b[i].ans = 1;
    }
    sort(b + 1, b + 1 + q, [&](Query &a, Query &b) {
        return a.id < b.id;
    });
    rep(i, 1, q) printf("%s
", b[i].ans? "TAK": "NIE");
}
/*
6 4 5 
1 5 2 6 
1 3 1 
2 3 2 
3 4 3 
4 5 5 
6 4 5 
4
1 2 4
2 6 9
1 5 9
6 5 8

*/