Ścieżki

W zadaniu mamy nieskierowany graf o $n$ wierzchołkach i $m$ krawędziach, w którym każdy wierzchołek jest pomalowany na jeden z $k$ ($k \leq 5$) kolorów. Należy znaleźć liczbę ścieżek prostych (czyli bez powtórzonych wierzchołków), w których wszystkie wierzchołki mają różne kolory.

Kluczowa obserwacja jest taka, że wystarczy ograniczyć się do ścieżek długości co najwyżej $k$ (bo każda dłuższa ścieżka prosta będzie zawierać powtórzenie koloru) oraz nie trzeba przejmować się tym, czy ścieżka jest prosta (bo ścieżki nie proste również zawierają powtórzenie koloru).

Ścieżki o długościach co najwyżej 4

Zliczyć ścieżki o długości 2 jest bardzo prosto -- są to wszystkie krawędzie, które mają końce różnych kolorów (przy czym każdą krawędź liczymy dwukrotnie, bo istotny jest również kierunek, w którym przechodzimy ścieżkę). Zatem czas zliczania tych ścieżek to $O(m)$.

Dla ścieżek długości 3 postaci $a-b-c$ możemy przeiterować się po wszystkich wyborach wierzchołka $b$, a następnie dobrać wierzchołki $a$ i $c$. Dobierać będziemy jedynie spośród sąsiadów wierzchołka $b$. Co więcej, interesować nas będzie tylko liczność sąsiadów w danym kolorze. Jeśli przez $Neigh[v]$ oznaczymy liczbę sąsiadów wierzchołka $v$, a przez $neigh[v][k]$ liczbę sąsiadów koloru $k$, to możemy obliczyć podwójną sumę (gdzie $k$ jest kolorem wierzchołka $v$), wybierającą kolor wierzchołka $a$, a następnie inny kolor dla wierzchołka $c$: $$\sum_{k_1 \neq k} \sum_{k_2 \not\in \{k_1,k\}} neigh[v][k_1] \cdot neigh[v][k_2]$$ albo równoważnie najpierw policzyć wszystkie kolorowania wierzchołków $a$ i $c$ kolorami różnymi od $k$, a następnie odjąć kolorowania, w których mają ten sam kolor: $$(Neigh[v] - neigh[v][k])^2 - \sum_{k_1 \neq k} (neigh[v][k_1])^2.$$ Złożoność czasowa będzie wynosić $O(m + nk)$.

Dla ścieżek długości 4 postaci $a-b-c-d$ robimy bardzo podobnie, tylko najpierw wybieramy krawędź $b-c$, a następnie dobieramy do niej wierzchołki $a$ i $d$. Złożoność czasowa będzie $O(mk)$.

Zauważmy, że powyższe rozwiązania działałyby również, gdybyśmy mieli znajdować ścieżki proste określonej długości, ale liczba różnych kolorów w grafie byłaby większa niż długości ścieżek.

Ścieżki długości 5

Dla ścieżek długości 5 rozwiązanie jest trochę bardziej skomplikowane i istotnie korzysta z tego, że ścieżka musi zawierać wszystkie kolory z grafu. Dla ścieżki $a-b-c-d-e$ najpierw iterujemy się po wierzchołku $c$, a następnie rozdzielamy pozostałe kolory pomiędzy podścieżki $c-b-a$ i $c-d-e$ (można to zrobić iterując się po jednym kolorze z podścieżki $c-b-a$ i zakładając, że drugi kolor jest minimalnym niewybranym, a następnie pozostałe kolory przydzielić do podścieżki $c-d-e$).

Dodatkowo, wcześniej spamiętujemy sobie liczbę kolorowań podścieżki typu $c-b-a$, jeśli wierzchołki $b$, $a$ mają kolory ze zbioru $\{k_1,k_2\}$. To można zrobić, iterując się po krawędzi $c-b$ i dobierając wierzchołek $a$.

Poniżej kod, który implementuje wszystkie przypadki:

const int N = 310000, K = 5;
int n,m,k,col[N],Neigh[N],neigh[N][K];
int edge[N][1<<K];
vector<int> adj[N];

ll sq(int x) { return ll(x) * x; }

int main() {
  scanf("%d%d%d",&n,&m,&k);
  REP(i,n) {
    scanf("%d",&col[i]);
    --col[i];
  }
  REP(i,m) {
    int a,b; scanf("%d%d",&a,&b);
    --a; --b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }

  ll ans = 0;

  // 2
  REP(i,n) for (int j : adj[i]) {
    if (col[i] != col[j]) ++ans;
  }

  // 3
  REP(i,n) for (int j : adj[i]) {
    neigh[i][col[j]]++;
    Neigh[i]++;
  }
  REP(i,n) {
    ans += sq(Neigh[i] - neigh[i][col[i]]);
    REP(k,K) if (k != col[i]) ans -= sq(neigh[i][k]);
  }

  // 4
  REP(i,n) for (int j : adj[i]) {
    if (col[i] != col[j]) {
      int ci = Neigh[i] - neigh[i][col[i]] - neigh[i][col[j]];
      int cj = Neigh[j] - neigh[j][col[j]] - neigh[j][col[i]];
      ans += ll(ci) * cj;
      REP(k,K) if (k != col[i] && k != col[j]) {
        ans -= ll(neigh[i][k]) * neigh[j][k];
      }
    }
  }

  // 5
  REP(i,n) for (int j : adj[i]) {
    if (col[i] != col[j]) {
      REP(k,K) if (k != col[i] && k != col[j]) {
        edge[i][(1<<col[j]) | (1<<k)] += neigh[j][k];
      }
    }
  }
  REP(i,n) {
    int k1 = col[i] == 0 ? 1 : 0;
    REP(k2,K) if (k2 != col[i] && k2 != k1) {
      int mask1 = (1<<k1) | (1<<k2);
      int mask2 = (1<<5)-1 - (1<<col[i]) - mask1;
      ans += 2 * ll(edge[i][mask1]) * edge[i][mask2];
    }
  }

  printf("%lld\n", ans);
}