competitive programming

Codeforces 869C The Intriguing Obsession

At first I got stuck with this problem.So I considered a simple case when there are only islands of two colors(c = 0).And it is easy to calculate the number of different ways under the constraints.

Suppose we only have islands with red and blue. Since we can only build bridges between a red island and a blue island, and further, we can only connect one blue island with one specific red island (otherwise distance between two blue islands is 2). And vice versa. Let f(i) be the number of different ways to build bridges when there are i bridges between red islands and blue islands.

Then we have:
$$
f(i) = C_a^i \cdot C_b^i \cdot i!
$$
Since we can pick i islands from red ones and i islands from blue ones, and we permutate i islands to match them.And the total number of it is:
$$
\sum_i^{min(a,b)} f(i)
$$

And next, we need to do with three colors. Try to add bridges between red islands and purple ones, and between blue and purple. And these bridges are exactly the same thing we calculated above. More importantly, bridges between red and purple, bridges between blue and purple are independent! And so is bridges between red and blue. The three groups of bridges are independent from each other. Change bridges in one group will not break the constraints. So our final answer is multiply the three number together.

Here is my code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

#include <iostream>
#include <algorithm>
#include <string.h>
#include <math.h>
#include <vector>

using namespace std; 

#define pb push_back
#define mp make_pair
#define ll long long
#define ull unsigned ll
#define db double
#define INF 0x3f3f3f3f
#define MOD 1000000007
#define PII pair<int, int>

#define mod 998244353

const int N=5010;
ll p[N],inv[N];
ll C(int n,int m) {
    return (((p[n]*inv[m])%mod)*inv[n-m])%mod;
}

ll qp(ll x,ll k) {
    ll ret=1;
    while (k>0) {
        if (k&1) ret=(ret*x)%mod;
        k>>=1;
        x=(x*x)%mod;
    }
    return ret;
}

void init() {
    p[0]=1;inv[0]=1;
    for (int i=1;i<=5000;i++) {
        p[i]=(p[i-1]*i)%mod;
        inv[i]=qp(p[i],mod-2);
    }
}

int a[3];

int main()
{
    init();
    scanf("%d%d%d",a,a+1,a+2);
    ll ans[3]={0};
    for (int i=0;i<3;i++) {
        int x=i,y=(i+1)%3;
        int lim=min(a[x],a[y]);
        for (int j=0;j<=lim;j++) {
            ll tmp=(((C(a[x],j)*C(a[y],j))%mod)*p[j])%mod;
            ans[i]=(ans[i]+tmp)%mod;
        }
    }
    ll res=(((ans[0]*ans[1])%mod)*ans[2])%mod;
    printf("%lld\n",res);
}

I use multiplicative inverse to calculate compositions of big number, the editorial uses dynammic programming instead.