Codeforces Round #439 (Div. 2) 869 C. The Intriguing Obsession

题目链接:C. The Intriguing Obsession

题目大意

三种颜色的点, 数量分别为a, b, c, 建边, 如果点i和点j有一条边直接相连, 那么ij的距离是1, 要求同种颜色的点距离必须大于等于3
求建边的方案总数(modulo 998 244 353)

思路

因为要求同色点距离大于等于3, 所以同色点之间不能建边
考虑ab两种颜色的点, 如果a色点连b色点, 这个b色点再连一个a色点, 那么就出现了一个同色点距离为2, 不符合题意, 所以两种颜色的点只能一一对应地连接
所以这ab两色之间建边方案为$ans_1 = \sum^{min(a, b)}_{k=0}{C^{k}_{a}C^{k}_{b}A^{k}_{k}}$
同理可得ac两色和bc两色之间建边方案数$ans_2, ans_3$
总的方案书数是$ans_1 \cdot ans_2 \cdot ans_3$

代码

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#include <set>
#include <queue>
#include <vector>
using namespace std;
typedef long long ll;
const int maxn = 5000+100, mod =  998244353;
ll a, b, c, f[maxn], p[maxn][maxn];
int main()
{
    f[0] = 1;
    for(int i=1; i<=5000; ++i) f[i] = (f[i-1] * i)%mod;
    memset(p,0,sizeof(p));
    p[0][0]=1;
    for (int i=1;i<=5000;i++)
    {
        p[i][0]=1;
        for (int j=1;j<=i;j++)
            p[i][j]=(p[i-1][j-1]+p[i-1][j])%mod;
    }

    ll ans[3] = {0, 0, 0};
    cin >> a >> b >> c;
    for(ll k=0; k<=min(a, b); ++k) ans[0] = (ans[0] + p[a][k] * p[b][k] % mod * f[k] % mod) % mod;
    for(ll k=0; k<=min(a, c); ++k) ans[1] = (ans[1] + p[a][k] * p[c][k] % mod * f[k] % mod) % mod;
    for(ll k=0; k<=min(c, b); ++k) ans[2] = (ans[2] + p[c][k] * p[b][k] % mod * f[k] % mod) % mod;
    cout << (ans[0]*ans[1] % mod * ans[2])%mod << endl;

    return 0;
}