牛客练习赛64D宝石装箱 (容斥+二维DP)

链接：https://ac.nowcoder.com/acm/contest/5633/D
来源：牛客网

时间限制：C/C++ 1秒，其他语言2秒
空间限制：C/C++ 262144K，其他语言524288K
64bit IO Format: %lld

题目描述

n颗宝石装进n个箱子使得每个箱子中都有一颗宝石。第i颗宝石不能装入第ai​个箱子。求合法的装箱方案对998244353取模。
两种装箱方案不同当且仅当两种方案中存在一颗编号相同的宝石装在不同编号的箱子中。

输入描述:

第一行一个整数n。
第二行n个整数$a_1,a_2,...,a_{n-1},a_n$​。

输出描述:

输出一行一个整数表示答案。

示例1

输入

复制

2
1 2

输出

复制

1

说明

只有一种合法的装箱方案2,1。

示例2

输入

复制

2
1 1

输出

复制

0

说明

没有合法的方案。

示例3

输入

复制

8
7 3 3 2 5 4 1 8

输出

复制

14568

备注:

1<=ai<=n<=8000{1<=a_i<=n<=8000}1<=ai​<=n<=8000

分析：

蛮有趣的一道容斥题,算是错排问题的拓展
关于错排可以看这里

为了方便分析，假设我们有a,b,c三个箱子，称每个箱子装入非法的球时的情况分别为A,B,C，每个箱子不能装入的球数分别为cntA,cntB,cntC，根据容斥原理有：
$A\cup B\cup C=A+B+C-A\cap B-B\cap C-A\cap C+A\cap B\cap C$
而题目要求合法方案数，根据集合运算性质有$\overline{A}\cap \overline{B}\cap \overline{C}=全集U-A\cup B\cup C$
变式得$\overline{A}\cap \overline{B}\cap \overline{C}=全集U-A-B-C+A\cap B+B\cap C+A\cap C-A\cap B\cap C$
推导得以下关系

已加入的盒子	1个盒子非法	2个盒子非法	3个盒子非法
A	$A$
A,B	$A+B$	$A\cap B$
A,B,C	$A+B+C$	$A\cap B+(A+B)\cap C$	$A\cap B\cap C$

将上述关系归纳,得到一个建立在集合上的DP数组公式：
假设我们新加入一个盒子K,前一个加入的盒子为J，
则 D P ( { A , . . . , J , K } , i ) = D P ( { A , . . . , J } , i ) + D P ( { A , . . . , J } , i − 1 ) ∩ K DP(\{A,...,J,K\},i)=DP(\{A,...,J\},i)+DP(\{A,...,J\},i-1)\cap K DP({A,...,J,K},i)=DP({A,...,J},i)+DP({A,...,J},i−1)∩K
我们用数组d(已统计的箱子数，非法箱子数)代表DP所对应的方案数，则
dp(0,0)=全集=$n!$

dp	1	2	3
1	$cntA\ast (n-1)!$
2	$(cntA+cntB)\ast (n-1)!$	$cntA\ast cntB\ast (n-2)!$
3	$(cntA+cntB+cntC)\ast (n-1)!$	$cntA\ast cntB\ast (n-2)!+(cntA+cntB)\ast cntC\ast (n-2)!$	$cntA\ast cntB\ast cntC\ast (n-3)!$

即$dp(k,i)=dp(k-1,i)+dp(k-1,i-1)\ast cnt(k)/(n-i)$
为了方便计算，我们把$cntA\ast cntB\ast cntC\ast ...$和$(n-i)!$分别用数组g和f来统计，再加上一点点滑动数组的思想，最终得到如下代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN=1e4;
ll cnt[MAXN];
ll g[MAXN];
ll f[MAXN];
const ll MOD=998244353;
ll mod_mul(ll a,ll b,ll n){
    ll res=0;
    while(b){
        if(b&1)res=(res+a)%n;
        a=(a+a)%n;
        b>>=1;
    }
    return res;
}
int main(){
	ios::sync_with_stdio(0);
	int n;
	cin>>n;
	g[0]=f[0]=1;
	for(int i=1;i<=n;i++){
		int x;
		cin>>x;
		cnt[x]++;
		f[i]=mod_mul(f[i-1],i,MOD);
	}
	for(int i=1;i<=n;i++){
		for(int j=i;j>0;j--){
			g[j]=(g[j]+mod_mul(g[j-1],cnt[i],MOD))%MOD;
		}
	}
	ll ans=0;
	for(int i=0,sig=1;i<=n;i++,sig*=-1){
		ans=(ans+sig*mod_mul(g[i],f[n-i],MOD)+MOD)%MOD;
	}
	cout<<ans;
}