제출 #1347070

#제출 시각아이디문제언어결과실행 시간메모리
1347070killerzaluuPaths (BOI18_paths)C++20
23 / 100
104 ms61860 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;
    e.reserve(m);

    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][6];
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= 5; j++) mp[i][j] = 0;
    }

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

    long long ans = 0;

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

    // length 3
    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;
    }

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

        int a[2], p = 0;
        for (int col = 1; col <= k; col++) {
            if (col != c[x] && col != c[y]) {
                a[p++] = col;
            }
        }

        if (p == 2) {
            ans += 2LL * (mp[x][a[0]] * mp[y][a[1]] + mp[x][a[1]] * mp[y][a[0]]);
        }
    }

    // s[x][i][j] = number of paths with colors i - j - c[x], ending at x
    static long long s[100005][6][6];
    for (int x = 1; x <= n; x++) {
        for (int y : g[x]) {
            int j = c[y];
            if (j == c[x]) continue;
            for (int i = 1; i <= k; i++) {
                if (i == c[x] || i == j) continue;
                s[x][i][j] += mp[y][i];
            }
        }
    }

    // length 5
    for (int x = 1; x <= n; x++) {
        for (int a = 1; a <= k; a++) {
            if (a == c[x]) continue;
            for (int b = 1; b <= k; b++) {
                if (b == c[x] || b == a) continue;
                for (int cc = 1; cc <= k; cc++) {
                    if (cc == c[x] || cc == a || cc == b) continue;
                    for (int d = 1; d <= k; d++) {
                        if (d == c[x] || d == a || d == b || d == cc) continue;
                        ans += s[x][a][b] * s[x][d][cc];
                    }
                }
            }
        }
    }

    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...