986A - Fair - CodeForces Solution

#include "bits/stdc++.h"
#define ll long long
#define fi first
#define se second
#define all(a) a.begin(), a.end()
#define rev reverse
#define pii pair<int, int>
#define vi vector<int>
#define vll vector<ll>
#define pll pair<ll, ll>
#define um unordered_map
#define us unordered_set
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define pb push_back
using namespace std;
using namespace __gnu_pbds;

typedef tree<
    int,
    null_type,
    less_equal<int>,
    rb_tree_tag,
    tree_order_statistics_node_update>
    ordered_multiset;
    

void solve() {
    int n, m, k, s; cin >> n >> m >> k >> s;
    vector<int> arr(n);
    for(int i=0; i<n; ++i) cin >> arr[i];
    vector<int> g[n];
    for(int i=0; i<m; ++i) {
        int a, b; cin >> a >> b;
        --a; --b;
        g[a].pb(b);
        g[b].pb(a);
    }
    ll dist[n][k+1];
    for(int i=0; i<n; ++i) {
        for(int j=1; j<=k; ++j) dist[i][j] = 1e7;
    }
    for(int i=1; i<=k; ++i) {
        queue<int> q;
        for(int j=0; j<n; ++j) {
            if(arr[j] == i) {
                q.push(j);
                dist[j][i] = 0;
            }
        }
        while(!q.empty()) {
            auto u = q.front();
            q.pop();
            for(auto& it: g[u]) {
                if(dist[it][i] > 1 + dist[u][i]) {
                    dist[it][i] = 1 + dist[u][i];
                    q.push(it);
                }
            }
        }
    }
    for(int i=0; i<n; ++i) {
        sort(dist[i]+1, dist[i] + k + 1);
        ll ans = 0;
        for(int j=1; j<=s; ++j) ans += dist[i][j];
        cout << ans << " ";
    }
}
int main() {
    ios_base :: sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    //cin >> t;
    while(t--) {
        solve();
        cout << "\n";
    }
    return 0;
}