けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

JOI 2011 春合宿 day2-2 Keycards(カードキー)

包除原理と聞いて

問題へのリンク

問題概要

JOI の春合宿が行われる施設の宿泊棟の部屋の鍵は,カードに穴がいくつか開いた形状をしている.穴を開ける位置の候補は N 個あり,これらのうちいくつかに穴を開けた 2N 枚の異なる鍵が作られた.

あなたは,JOI の春合宿のための 1 枚以上 2N 枚以下の鍵をまとめて受け取った.穴を開ける候補の位置を合わせて鍵を重ねてみたあなたは,ちょうど K か所,受け取ったすべての鍵に穴が開いている位置があることに気づいた.

このようなことが起こるような,受け取った鍵の組み合わせは何通りあるだろうか.答えを 1 000 000 007 (素数) で割った余りを求めよ.

制約

  •  1 \le K \le N \le 1000000

考えたこと

ぱっと見、 K があまり意味ないように思えた。というのも

穴の開ける箇所を  C(N, K) 通り決めてしまったら、残りは  K = 0 についての問題に帰着できる。

つまりは「 N 箇所すべてが何かしらの鍵によって閉じられている」ような場合の数を求めれば良い。いかにも包除原理っぽい。

  •  1 箇所以上穴が空いている場合を引く
  •  2 箇所以上穴が空いている場合を足す
  • ...
#include <iostream>
#include <vector>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return true; } return false; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return true; } return false; }

const int MAX = 2100000;
const int MOD = 1000000007;

long long fac[MAX], finv[MAX], inv[MAX];
void COMinit(){
    fac[0] = fac[1] = 1;
    finv[0] = finv[1] = 1;
    inv[1] = 1;
    for(int i = 2; i < MAX; i++){
        fac[i] = fac[i-1] * i % MOD;
        inv[i] = MOD - inv[MOD%i] * (MOD/i) % MOD;
        finv[i] = finv[i-1] * inv[i] % MOD;
    }
}
long long COM(int n, int k){
    if(n < k) return 0;
    if (n < 0 || k < 0) return 0;
    return fac[n] * (finv[k] * finv[n-k] % MOD) % MOD;
}

int main () {
    COMinit();
    int N, K; cin >> N >> K;
    long long factor = COM(N, K);
    N -= K;
    
    if (N == 0) cout << factor << endl;
    else {
        // 2 のべき乗
        vector<long long> twobeki(N+1);
        twobeki[0] = 2;
        for (int i = 0; i < N; ++i) twobeki[i + 1] = (twobeki[i] * twobeki[i]) % MOD;
    
        // i 個以上空いているところ
        long long res = twobeki[N];
        for (int i = 1; i <= N; ++i) {
            long long tmp = COM(N, i) * twobeki[N-i] % MOD;
            if (i & 1) res = (res - tmp + MOD) % MOD;
            else res = (res + tmp) % MOD;
        }
        cout << res * factor % MOD << endl;
    }
}