AT_agc025_b RGB Coloring题解

题目链接

集训讲了这道题，于是写个题解巩固一下。

思路

题意转化为数学公式，就是在求

\[\sum_{x=0}^n\sum_{y=0}^n\sum_{z=0}^n\left [ Ax+By+(A+B)z=k\wedge x+y+z\le n\right ] \]

如果直接暴力枚举前两种组合，判断 \(k-(Ax+By)\) 是否为 \((A+B)\) 的正整数倍，时间复杂度为 \(O(n^2)\)，无法通过本题。

重新再看看有什么发现，不难看出，这个 \(A+B\) 很独特，出题人给定 \(A+B\) 一定因其有一些独特之处。考虑从这里切入。

不难发现，其实 \(A+B\) 就相当于红蓝两种颜色同时涂在一个格子上，效果叠加而已。因此转化后每一种颜色就可以独立看待，求

\[\sum_{x=0}^n\sum_{y=0}^n\left[Ax+By=k\right] \]

这样就简单了，直接枚举 \(x\)，判断 \(k-Ax\) 是否为 \(B\) 的正整数倍即可。

线性递推求逆元和阶乘时间复杂度降为 \(O(n)\)。
如果用费马小定理求逆元，时间复杂度为 \(O(n\log n)\)，也可以通过本题。

注意开 long long。

代码

#include<iostream>
using namespace std;
const int N=3e5+10,mod=998244353;
typedef long long ll;
ll n,a,b,k,res;
ll fac[N],inv[N];
void init()//预处理阶乘和逆元
{
    fac[0]=fac[1]=inv[0]=inv[1]=1;
    for(int i=2;i>=n;i++)
    {
        fac[i]=fac[i-1]*i%mod;
        inv[i]=(mod-mod/i)*inv[mod%i]%mod;
    }
    for(int i=2;i<=n;i++)inv[i]=inv[i]*inv[i-1]%mod;
}
ll C(ll n,ll m)
{
	if(n<m)return 0;
	return fac[n]*inv[m]%mod*inv[n-m]%mod;
}
int main()
{
	cin>>n>>a>>b>>k;
    init();
	for(ll i=0;i*a<=k&&i<=n;i++)
		if((k-a*i)%b==0)
			res=(res+C(n,i)*C(n,(k-a*i)/b)%mod)%mod;
	cout<<res;
	return 0;
}

不要直接 ctrl c 哦。