Submission #1347037

#TimeUsernameProblemLanguageResultExecution timeMemory
1347037killerzaluuPaths (BOI18_paths)C++20
70 / 100
128 ms31664 KiB
#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, m, k;
    cin >> n >> m >> k;

    int c[n + 1];
    for (int i = 1; i <= n; i++) cin >> c[i];

    vector<int> g[n + 1];
    vector<pair<int,int>> e;

    for (int i = 1; i <= m; i++) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
        e.push_back({a, b});
    }

    long long mp[n + 1][5];
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= 4; j++) mp[i][j] = 0;
    }

    for (int v = 1; v <= n; v++) {
        for (int u : g[v]) {
            mp[v][c[u]]++;
        }
    }

    long long ans = 0;

    for (auto [x, y] : e) {
        if (c[x] != c[y]) ans += 2;
    }

    for (int x = 1; x <= n; x++) {
        long long s = 0, sq = 0;
        for (int col = 1; col <= k; col++) {
            if (col == c[x]) continue;
            s += mp[x][col];
            sq += mp[x][col] * mp[x][col];
        }
        ans += s * s - sq;
    }

    for (auto [x, y] : e) {
        if (c[x] == c[y]) continue;

        int a = 0, b = 0;
        for (int col = 1; col <= 4; col++) {
            if (col != c[x] && col != c[y]) {
                if (a == 0) a = col;
                else b = col;
            }
        }

        ans += 2LL * (mp[x][a] * mp[y][b] + mp[x][b] * mp[y][a]);
    }

    cout << ans << '\n';
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...