CodeForces 986A - Fair (BFS)

思路：BFS从多源同时求each货物到each地点的最短距离

void nth_element(beg, nth, end) : 求前s个dist的和。

#include <bits/stdc++.h>

using namespace std;
const int maxn = 1e5 + 5;
const int inf = 1e7 + 5;

vector<int> g[maxn], forcol[105];

int n, m, k, s;
int dist[maxn];
int ans[maxn][105];

void dfs(int c) {
    for(int i = 1; i <= n; ++i) dist[i] = inf;
    queue<int> q;
    for(int i : forcol[c]) {
        dist[i] = 0;
        q.push(i);
    }
    while(q.size()) {
        int u = q.front(); q.pop();
        for(int v : g[u]) {
            if(dist[v] <= dist[u] + 1) continue;
            dist[v] = dist[u] + 1;
            q.push(v);
        }
    }
    for(int i = 1; i <= n; ++i)  ans[i][c] = dist[i];
}

int main() {
    cin >> n >> m >> k >> s;
    for(int i = 1, x; i <= n; ++i) {
        scanf("%d", &x);
        forcol[x].push_back(i);
    }
    for(int i = 0; i < m; ++i) {
        int u, v;
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for(int col = 1; col <= k; ++col) {
        dfs(col);
    }
    for(int i = 1; i <= n; ++i) {
        nth_element(ans[i] + 1, ans[i] + s + 1, ans[i] + k + 1);
        int res = 0;
        for(int j = 1; j <= s; ++j) {
            res += ans[i][j];
        }
        if(i > 1) putchar(' ');
        printf("%d", res);
    }
    cout << endl;
    return 0;
}