[ABC266G] Yet Another RGB Sequence

Problem Statement

You are given integers $R$, $G$, $B$, and $K$. How many strings $S$ consisting of R, G, and B satisfy all of the conditions below? Find the count modulo $998244353$.

• The number of occurrences of R, G, and B in $S$ are $R$, $G$, and $B$, respectively.
• The number of occurrences of RG as (contiguous) substrings in $S$ is $K$.

Constraints

• $1 \leq R,G,B\leq 10^6$
• $0 \leq K \leq \mathrm{min}(R,G)$
• All values in input are integers.

---

Input

Input is given from Standard Input in the following format:

$R$ $G$ $B$ $K$

Output

Print the answer.

---

Sample Input 1

2 1 1 1

Sample Output 1

6

The following six strings satisfy the conditions.

• RRGB
• RGRB
• RGBR
• RBRG
• BRRG
• BRGR

---

Sample Input 2

1000000 1000000 1000000 1000000

Sample Output 2

80957240

Find the count modulo $998244353$.

这个数据范围一看就知道不好 dp，大概率是推式子。

恰好 \(k\) 个不好求，但是我们可以想到一个答案很接近的方法：将一个 RG 打包成一个数，然后让其参与排列。

按照上面这种方法，定义 \(f(i)\) 为如果有 \(i\) 个 RG，将其打包成一份后算出来的答案。易得 \(f(i)=C_{r+g+b-i}^{b}*C_{r+g-i}^{i}\)

如果定义 \(g(i)\) 为恰好有 \(i\) 个 RG 的方案数，那么考虑一个 \(f(i)\) 中算了几次 \(g(j)\)，那么会发现对于一种方案中 \(j\) 个 RG，选出 \(i\) 个 RG 的方案都会被算一次。那么得到 \(f(i)=\sum\limits_{j=i}^{\infin}C_{j}^ig(j)\)

上二项式反演公式， \(g(i)=\sum\limits_{j=i}^{\infin}(-1)^{j-i}C_{j}^if(j)\)

那么这个问题就解决了。

#include<cstdio>
const int N=3e6+5,P=998244353;
int jc[N],r,g,b,k,n,ans,inv[N];
int pown(int x,int y)
{
	if(!y)
		return 1;
	int t=pown(x,y>>1);
	if(y&1)
		return 1LL*t*t%P*x%P;
	return 1LL*t*t%P;
}
int ct(int x,int y)
{
	if(x<y)
		return 0;
	return 1LL*jc[x]*inv[y]%P*inv[x-y]%P;
}
int calc(int x)
{
	n=r+g+b-x;
	return 1LL*ct(n,b)%P*ct(n-b,x)%P*ct(n-b-x,r-x)%P;
}
int main()
{
	scanf("%d%d%d%d",&r,&g,&b,&k),inv[0]=1;
	for(int i=jc[0]=1;i<=r+g+b;i++)
		jc[i]=1LL*jc[i-1]*i%P,inv[i]=pown(jc[i],P-2);
	for(int i=k;i<=r&&i<=g;i++)
		ans+=((i-k&1)? -1LL:1LL)*ct(x,k)*calc(i)%P,ans=(1LL*ans+P)%P;
	printf("%d",ans);
}