読者です 読者をやめる 読者になる 読者になる

ゆらのふなびと

競プロ, Python, C++

yukicoder No.389 ロジックパズルの組み合わせ

問題

M個のマスが1直線状に並んでいる。K個のヒントがあり、k番目のヒントは連続するH_k個のマスを塗れというものである。ただし異なるヒントによって塗られるマスどうしは隣接してはならない。このような塗り方は何通りあるか。mod 109 + 7 で答えよ。ただし、すべてのヒントを守るような塗り方が存在しない場合は "NA" と出力せよ。

制約

1 <= M <= 105

ΣH_k <= M

H_kは非負整数。ただし0が現れるのは K = 1 のときのみ(このとき答えは「どのマスも塗らない」の1通りとする)。

解法

問題文の通り、ヒントが0のみのときはコーナーケース。以下ではヒントが0を含まないときを考える。

ヒントの指示は「H_1, H_2, ..., H_k 個の連続するマスを隣り合わないように塗れ」ということだが、これは言い換えると「白いマスの間または両端の M - ΣH_k + 1 個の位置から K 個選んで、左から順に長さ H_k の黒いマスを入れろ」という意味。位置の組み合わせさえ決まれば黒いマスは左から順に入れるという1通りしかないので、答えは (M - ΣH_k + 1)_C_K 通り。ただし M - ΣH_k + 1 < K のときは K 個の位置を選べないので "NA" 。

組み合わせを求めるには逆元を使った様々な方法があるが、以下のコードでは定義に従って愚直に計算する方法を用いた。計算量は O(K)。

#include <bits/stdc++.h>
using namespace std;
#define int long long   // <-----!!!!!!!!!!!!!!!!!!!

#define rep(i,n) for (int i=0;i<(n);i++)
#define rep2(i,a,b) for (int i=(a);i<(b);i++)
#define rrep(i,n) for (int i=(n)-1;i>=0;i--)
#define rrep2(i,a,b) for (int i=(b)-1;i>=(a);i--)
#define all(a) (a).begin(),(a).end()

typedef long long ll;
typedef pair<int, int> Pii;
typedef tuple<int, int, int> TUPLE;
typedef vector<int> V;
typedef vector<V> VV;
typedef vector<VV> VVV;
typedef vector<vector<int>> Graph;
const int inf = 1e9;
const int mod = 1e9 + 7;

ll powmod (ll a, ll p) {
    ll ans = 1;
    ll mul = a;
    for (; p > 0; p >>= 1, mul = (mul * mul) % mod) {
        if ((p & 1) == 1) ans = (ans * mul) % mod;
    }
    return ans;
}

ll getRev(ll x) {
    return powmod(x, mod - 2);
}

ll C(ll n, ll k) {
    ll ret = 1;
    for (ll i = n; i > n - k; i--) {
        (ret *= i) %= mod;
    }
    for (ll i = k; i >= 1; i--) {
        (ret *= getRev(i)) %= mod;
    }
    return ret;
}

signed main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(0);

    int M;
    cin >> M;
    V h;
    string s;
    cin.ignore();
    getline(cin, s);
    stringstream ss(s);
    int x;
    while (ss >> x) h.emplace_back(x);
    int K = h.size();

    int sm = accumulate(all(h), 0);
    if (sm == 0) {
        cout << 1 << endl;
        return 0;
    }
    if (M - sm + 1 < K) {
        cout << "NA" << endl;
        return 0;
    }
    cout << C(M - sm + 1, K) << endl;
}

☆2で逆元ひえぇ……と思っていたら☆3になった。