Submission #441452

#	Time	Username	Problem	Language	Result	Execution time	Memory
441452		abc864197532	Paths (BOI18_paths)	C++17	100 / 100	265 ms	20288 KiB

paths

#include <bits/stdc++.h>
using namespace std;
#define lli long long int
#define X first
#define Y second
#define pb push_back
#define eb emplace_back
#define mp make_pair
#define pii pair<int, int>
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
#define test(args...) abc("[" + string(#args) + "]", args)
void abc() {cerr << endl;}
template <typename T, typename ...U> void abc(T a, U ...b) {
    cerr << a << ' ', abc(b...);
}
template <typename T> void printv(T l, T r) {
    while (l != r) cout << *l << " \n"[++l == r];
}
const int mod = 1e9 + 7, N = 100000;

int cnt2[N][5][5];

int main () {
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n, m, k;
    cin >> n >> m >> k;
    vector <int> color(n, 0);
    vector <pii> edge;
    vector <vector <int>> cnt(n, vector <int>(k, 0));
    for (int i = 0; i < n; ++i) cin >> color[i], --color[i];
    for (int i = 0, u, v; i < m; ++i) {
        cin >> u >> v, --u, --v;
        if (color[u] == color[v]) continue;
        edge.eb(u, v);
        cnt[u][color[v]]++;
        cnt[v][color[u]]++;
    }
    lli ans = 0;
    // len = 2
    ans += edge.size() * 2;
    // len = 3;
    for (pii A : edge) {
        int u, v;
        tie(u, v) = A;
        for (int i = 0; i < k; ++i) if (i != color[v] && i != color[u]) {
            ans += cnt[u][i] + cnt[v][i];
        }
    }
    // len = 4
    for (pii A : edge) {
        int u, v;
        tie(u, v) = A;
        for (int i = 0; i < k; ++i) for (int j = 0; j < k; ++j) {
            if (i == j) continue;
            if (i == color[v] || i == color[u]) continue;
            if (j == color[v] || j == color[u]) continue;
            ans += 1ll * cnt[v][i] * cnt[u][j] + 1ll * cnt[v][j] * cnt[u][i];
        }
    }
    if (k <= 4) return cout << ans << endl, 0;
    // len = 5
    for (pii A : edge) {
        int u, v;
        tie(u, v) = A;
        for (int t : {0, 1}) {
            for (int i = 0; i < k; ++i) if (i != color[v]) {
                cnt2[u][color[v]][i] += cnt[v][i];
            }
            swap(u, v);
        }
    }
    for (pii A : edge) {
        int u, v;
        tie(u, v) = A;
        for (int i = 0; i < k; ++i) if (i != color[u] && i != color[v]) {
            int xorr = 4 ^ i ^ color[u] ^ color[v];
            for (int j = 0; j < k; ++j) if (j != i && j != color[u] && j != color[v]) {
                int nj = xorr ^ j;
                ans += 1ll * cnt[u][i] * cnt2[v][j][nj] + 1ll * cnt[v][i] * cnt2[u][j][nj];
            }
        }
    }
    cout << ans << endl;
}

Compilation message (stderr)

paths.cpp: In function 'int main()':
paths.cpp:67:18: warning: unused variable 't' [-Wunused-variable]
   67 |         for (int t : {0, 1}) {
      |                  ^

• Subtask 1: 23 / 23
• Subtask 2: 20 / 20
• Subtask 3: 27 / 27
• Subtask 4: 30 / 30