yukicoder : No.1691 Badugi

想定解それなのね…。
https://yukicoder.me/problems/no/1691

問題

NM枚のカードがある。
1～N番までのN通りの色と、1～M番までのM通りの番号で、NM枚のカード中に色と番号の組み合わせは全通り登場する。
ここからK枚カードを選ぶとき、以下を満たすのは何通りか。

• K枚中(K-2)枚のカード選ぶとき、全カードで色も番号も異なるような組み合わせがある。

解法

K枚選んだ時、色か番号がいずれも(K-2)種類以上であり、かつ片方がピッタリ(K-2)種類であれば、条件を満たす。
まず(K-2)枚、色も番号も異なるカードを選んだとする。これはC(N,K-2)*P(M,K-2)通りである。
後2枚カードを選んだ時、条件を満たす(K-2)毎の組み合わせが何通り得られるかを考える。
もしn通り得られるような組み合わせは、n重にカウントしてしまうのでnで割ればよい。

最初の(K-2)枚に対し、残り2枚を以下から選ぶ場合を考える。

• A: 色も番号も、それぞれ最初の(K-2)枚中のいずれかと一致する。
• B: 色は最初の(K-2)枚中のいずれかと一致するが、番号は一致しない
• C: 番号は最初の(K-2)枚中のいずれかと一致するが、色は一致しない

2枚の選び方が(A,A),(A,B),(A,C),(B,B),(B,C),(C,C)のそれぞれの場合、何重にカウントされるかを細かく条件分けして数え上げて行こう。
具体的に何重かはEditorialを参照。

ll N,M,K;
const ll mo=998244353;
const int NUM_=1400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll comb(ll N_, ll C_) {
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(ll P_,ll Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	ll ret=0;
	ll base=comb(N,K-2)*comb(M,K-2)%mo;
	base=base*fact[K-2]%mo;
	
	// AA
	ll rect=1LL*(K-2)*(K-3)/2%mo;
	ret+=rect*modpow(2)%mo;
	ll lef=1LL*(K-3)*(K-2)%mo;
	lef=(mo+lef*(lef-1)/2%mo-rect)%mo;
	ret+=lef;
	
	//AB
	ll b=(N-(K-2))*(K-2)%mo;
	ret+=b%mo*(K-3)%mo*modpow(3)%mo;
	ret+=b%mo*(K-3)%mo*(K-3)%mo*modpow(2)%mo;
	// BB
	ll tot=b*(b-1)%mo*modpow(2)%mo;
	ll cr=((N-(K-2))*(N-(K-2)-1)/2)%mo*(((K-2)*(K-3))%mo)%mo;
	ret+=(tot-cr)%mo*modpow(3)%mo;
	ret+=cr%mo*modpow(4)%mo;
	// AC
	ll c=(M-(K-2))*(K-2)%mo;
	ret+=c%mo*(K-3)%mo*modpow(3)%mo;
	ret+=c%mo*(K-3)%mo*(K-3)%mo*modpow(2)%mo;
	
	// CC
	tot=c*(c-1)%mo*modpow(2)%mo;
	cr=((M-(K-2))*(M-(K-2)-1)/2)%mo*(((K-2)*(K-3))%mo)%mo;
	ret+=(tot-cr)%mo*modpow(3)%mo;
	ret+=cr%mo*modpow(4)%mo;
	// BC
	ret+=b%mo*(K-3)%mo*(M-(K-2))%mo*modpow(4)%mo;
	
	ret=(ret%mo+mo)*base%mo;
	cout<<ret%mo<<endl;
}

まとめ

本番中、この案もあるなと思ったものの、もっとスマートな方法がないかと探し回ってしまった。