12752번: Cities

https://www.acmicpc.net/problem/12752

#include<cstdio>
#include<algorithm>
#include<vector>
#include<queue>
using namespace std;
const int MXN = 1e5;
typedef long long ll;
int n, k, m, a[6], p[5];
ll dis[6][MXN + 2], r = 1e15, cst[5][5][MXN + 1];
vector<pair<int, ll> > adj[MXN + 2];
void spath(int h) {
    fill(dis[h], dis[h + 1], 1e15);
    dis[h][a[h]] = 0;
    priority_queue<pair<ll, int> > pq;
    pq.push({ 0,a[h] });
    while (!pq.empty()) {
        int tpos = pq.top().second;
        ll tdis = -pq.top().first;
        pq.pop();
        if (tdis^dis[h][tpos]) continue;
        for (auto it : adj[tpos]) if (it.second + dis[h][tpos] < dis[h][it.first]) {
            dis[h][it.first] = it.second + dis[h][tpos];
            pq.push({ -dis[h][it.first],it.first });
        }
    }
}
int main() {
    scanf("%d%d%d", &n, &k, &m);
    for (int i = 0; i < k; i++) scanf("%d", a + i), p[i] = i;
    for (int i = 0, x, y, z; i < m; i++) {
        scanf("%d%d%d", &x, &y, &z);
        adj[x].push_back({ y,z });
        adj[y].push_back({ x,z });
    }
    for (int i = 0; i < k; i++) spath(i);
    for (int i = 1; i <= n; i++) {
        adj[0].push_back({ i,0 });
        adj[i].push_back({ n + 1,0 });
        ll s = 0;
        for (int j = 0; j < k; j++) s += dis[j][i];
        r = min(r, s);
    }
    if (k > 3) {
        for (int i = 1; i < k; i++) {
            for (int j = 0; j < i; j++) {
                for (int u = 1; u <= n; u++) {
                    adj[0][u - 1].second = dis[i][u] + dis[j][u];
                    adj[u].rbegin()->second = 0;
                    for (int v = 0; v < k; v++) if (v^i&&v^j) adj[u].rbegin()->second += dis[v][u];
                }
                spath(k);
                r = min(r, dis[k][n + 1]);
                for (int u = 1; u <= n; u++) cst[i][j][u] = dis[k][u];
            }
        }
    }
    if (k == 5) {
        for (int i = 0; i < k; i++) {
            for (int p1 = 1; p1 < k; p1++) if (p1^i) {
                for (int p2 = 0; p2 < p1; p2++) if (p2^i) {
                    int q[2], qcnt = 0;
                    for (int u = 0; u < k; u++) if (u^i&&u^p1&&u^p2) q[qcnt++] = u;
                    for (int u = 1; u <= n; u++) r = min(r, cst[p1][p2][u] + cst[q[1]][q[0]][u] + dis[i][u]);
                }
            }
        }
    }
    printf("%lld", r);
    return 0;
}