제출 #1143512

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1143512		LucaLucaM	Paths (BOI18_paths)	C++20	100 / 100	160 ms	42808 KiB

paths

#include <iostream>
#include <vector>
#include <algorithm>
#include <cassert>

using ll = long long;
#define debug(x) #x << " = " << x << '\n'

const int NMAX = 3e5;

int a[NMAX + 1];
std::vector<int> g[NMAX + 1];

int n, k;

ll count1() {
  return n;
}
ll count2() {
  ll ret = 0;
  for (int i = 1; i <= n; i++) {
    for (const auto &j : g[i]) {
      if (a[i] != a[j]) {
        ret++;
      }
    }
  }
  return ret;
}
ll count3() {
  ll ret = 0;
  for (int i = 1; i <= n; i++) {
    std::vector<int> f(6, 0);
    for (const auto &j : g[i]) {
      f[a[j]]++;
    }
    for (int c1 = 1; c1 <= k; c1++) {
      for (int c2 = 1; c2 <= k; c2++) {
        if (c1 != c2 && a[i] != c1 && a[i] != c2) {
          ret += (ll) f[c1] * f[c2];
        }
      }
    }
  }
  return ret;
}

ll count4() {
  ll ret = 0;
  std::vector<std::vector<ll>> f(n + 1, std::vector<ll>(6, 0));
  for (int i = 1; i <= n; i++) {
    for (const auto &j : g[i]) {
      f[i][a[j]]++;
    }
  }
  for (int i = 1; i <= n; i++) {
    for (const auto &j : g[i]) { // O(m)
      if (a[i] != a[j]) {
        for (int c1 = 1; c1 <= 5; c1++) {
          if (c1 != a[i] && c1 != a[j]) {
            for (int c2 = 1; c2 <= 5; c2++) {
              if (c2 != a[i] && c2 != a[j] && c1 != c2) { 
                ret += (ll) f[i][c1] * f[j][c2];
              }
            }
          }
        }
      }
    }
  }

  return ret;
}

ll count5() {
  std::vector<std::vector<ll>> f(n + 1, std::vector<ll>(6, 0));
  for (int i = 1; i <= n; i++) {
    for (const auto &j : g[i]) {
      f[i][a[j]]++;
    }
  }

  ll ret = 0;
  for (int i = 1; i <= n; i++) {
    std::vector<std::vector<ll>> ff(6, std::vector<ll>(6, 0));
    for (const auto &j : g[i]) {
      if (a[i] != a[j]) {
        for (int c = 1; c <= 5; c++) {
          if (c != a[i] && c != a[j]) {
            ff[a[j]][c] += f[j][c];
          }
        }
      }
    }

    for (int c1 = 1; c1 <= 5; c1++) {
      if (a[i] != c1) {
        for (int c2 = 1; c2 <= 5; c2++) {
          if (c1 != c2 && a[i] != c2) {
            for (int c3 = 1; c3 <= 5; c3++) {
              if (c2 != c3 && c1 != c3 && a[i] != c3) {
                int c4 = (1 ^ 2 ^ 3 ^ 4 ^ 5 ^ c1 ^ c2 ^ c3 ^ a[i]);
                assert(c4 != c3 && c4 != c2 && c4 != c1 && c4 != a[i]);
                assert(c1 + c2 + c3 + c4 + a[i] == 15);
                assert((c1 ^ c2 ^ c3 ^ c4 ^ a[i]) == (1 ^ 2 ^ 3 ^ 4 ^ 5));
                ret += (ll) ff[c1][c2] * ff[c3][c4];
              }
            }
          }
        }
      }
    }
  }

  return ret;
}

int main() {
  #ifdef LOCAL
    freopen("input.txt", "r", stdin);
  #endif
  std::ios_base::sync_with_stdio(false);
  std::cin.tie(0);
  
  int m;
  std::cin >> n >> m >> k;

  for (int i = 1; i <= n; i++) {
    std::cin >> a[i];
  }

  while (m--) {
    int u, v;
    std::cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }

  if (k == 1) {
    std::cout << 0;
  } else if (k == 2) {
    std::cout << count2();
  } else if (k == 3) {
    std::cout << count2() + count3();
  } else if (k == 4) {
    std::cout << count2() + count3() + count4();
  } else {
    std::cout << count2() + count3() + count4() + count5();
  }

  return 0;
}

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