CF198 Div1 C - Iahub and Permutations

Codeforces

包除原理って良いなぁ😇

問題

Iahubは長さnの順列aを持っていたが、席を外している間に会社の同僚によってaのいくつかの要素を-1に書き換えられてしまった
書き換えられる前のaは、a[i]=iとなるような要素が1つもなかったことだけは覚えている
元のaとしてあり得るものの個数を求めよ

2 <= n <= 2000
aは-1になっている要素が少なくとも2つあり、条件を満たす順列が少なくとも1つは存在することが保証される

解法

-1の個数をk、消された数字xのうちa[x]が-1になっているようなものの個数をmとし、
f(i): a[x]=xとなるようなxがi個以上存在するような順列aの個数
とすると、包除原理より答えは
Σ f(i) * (-1)^i (0 <= i <= m)
となり、
f(i) = (m個からi個選ぶ選び方) * (残りk-i個を好きに当てはめる方法の個数)
= mCi * (k-i)!

#define int long long

int n, a[2020], b[2020], f[2020], m, k, res;

const int MAXS = 2020;
ll fact[MAXS+1], factr[MAXS+1], inv[MAXS+1];
ll comb(ll n, ll r) {
	if (fact[0]==0) {
		fact[0] = factr[0] = inv[1] = 1;
		for (int i=2; i<=MAXS; i++) inv[i] = inv[mod%i] * (mod-mod/i) %mod;
		for (int i=1; i<=MAXS; i++) fact[i] = fact[i-1]*i %mod, factr[i] = factr[i-1]*inv[i] %mod;
	}
	if (r<0 || n<r) return 0;
	return fact[n]*factr[r] %mod *factr[n-r] %mod;
}

signed main() {
	f[0] = 1;
	reps(i,1,2020) f[i] = f[i-1]*i %mod;
	cin >> n;
	rep(i,n) {
		cin >> a[i];
		if (a[i]!=-1) b[a[i]-1] = 1;
		else k++;
	}
	rep(i,n) if (!b[i] && a[i]==-1) m++;
	rep(i,m+1) (res += comb(m,i) * f[k-i] * (i%2 ? -1 : 1)) %= mod;
	(res += mod) %= mod;
	cout << res << endl;
}

これくらいの問題なら瞬殺できるようになってきた

http://codeforces.com/problemset/problem/340/E