LeeCode 3130 DP / 组合数学

题意

传送门 LeeCode 3130 找出所有稳定的二进制数组 II

题解

DP

令$dp(i,j,k)$表示长度为$i$的序列中1的数量为$j$，且序列以$k$结尾时的序列数量。那么只需要暴力枚举每一段连续且元素相同的序列长度，递推即可。此时时间复杂度$O((zero+one{)}^{2}one)$，难以胜任。可以通过维护前缀和，或者直接利用递推关系作差分，优化掉枚举序列长度这一维。如下采用维护前缀和进行优化，时间复杂度$O((zero+one)one)$。

#include <bits/stdc++.h>
using namespace std;

class Solution {
   public:
    int numberOfStableArrays(int zero, int one, int limit) {
        int n = zero + one;
        int dp[n + 1][one + 1][2];
        int diag_sum[n + 2][one + 2];
        int y_sum[one + 1][n + 2];
        memset(dp, 0, sizeof(dp));
        memset(diag_sum, 0, sizeof(diag_sum));
        memset(y_sum, 0, sizeof(y_sum));
        const int mod = 1e9 + 7;
        auto add = [&](int x, int y) {
            x += y;
            if (x >= mod) {
                x -= mod;
            }
            return x;
        };
        dp[0][0][0] = dp[0][0][1] = 1;
        diag_sum[1][1] = 1;
        y_sum[0][1] = 1;
        for (int i = 1; i <= n; ++i) {
            for (int j = 0; j <= one; ++j) {
                dp[i][j][0] = add(y_sum[j][i], mod - y_sum[j][max(0, i - limit)]);
                int d = min({i, j, limit});
                dp[i][j][1] = add(diag_sum[i][j], mod - diag_sum[i - d][j - d]);
                diag_sum[i + 1][j + 1] = add(dp[i][j][0], diag_sum[i][j]);
                y_sum[j][i + 1] = add(y_sum[j][i], dp[i][j][1]);
            }
        }
        int res = add(dp[n][one][0], dp[n][one][1]);
        return res;
    }
};

组合数学

可以通过枚举连续的0/1的子段的数量统计答案。具体而言，令0的段数量为$n$，因为0/1子段是交错的，故1的段数量仅有$n-1,n,n+1$三种可能。0/1子段可以分别独立地统计满足条件的序列数量，最后相乘即可。以0为例，假设其子段数量为$n$，则每一段至少需要分配一个0。故剩余可自由分配的0数量为$m=\text{zero}-n$。问题转化为$m$个元素分配至$n$个桶且每个桶的元素数量不大于等于$limit$的方案数，可以通过容斥原理$O(min(n,m/limit))$求解。不妨设$\text{zero}\le \text{one}$，总时间复杂度$O(zero\times \min (zero,one/limit))$。

#include <bits/stdc++.h>
using namespace std;

class Solution {
   public:
    int numberOfStableArrays(int zero, int one, int limit) {
        if (zero > one) {
            swap(zero, one);
        }
        const int mod = 1e9 + 7;
        vector<long long> fac(one + 1), inv(one + 1), invf(one + 1);
        fac[0] = fac[1] = 1;
        inv[1] = 1;
        invf[0] = invf[1] = 1;
        for (int i = 2; i <= one; ++i) {
            fac[i] = fac[i - 1] * i % mod;
            inv[i] = (mod - mod / i) * inv[mod % i] % mod;
            invf[i] = invf[i - 1] * inv[i] % mod;
        }
        auto get_comb = [&](int n, int m) -> long long {
            if (n < 0 || m < 0 || n < m) {
                return 0;
            }
            return fac[n] * invf[m] % mod * invf[n - m] % mod;
        };
        auto get = [&](int m, int n) -> long long {
            if (m < 0) {
                return 0;
            }
            long long res = 0;
            for (int k = 0; k <= n && m >= limit * k; ++k) {
                long long d = get_comb(n, k) * get_comb(m - limit * k + n - 1, n - 1) % mod;
                if (~k & 1) {
                    res += d;
                } else {
                    res -= d;
                }
            }
            res = (res % mod + mod) % mod;
            return res;
        };
        vector<long long> f_zero(zero + 1), f_one(zero + 2);
        for (int i = 1; i <= zero; ++i) {
            f_zero[i] = get(zero - i, i);
        }
        for (int i = 1; i <= zero + 1; ++i) {
            f_one[i] = get(one - i, i);
        }
        long long res = 0;
        for (int i = 1; i <= zero; ++i) {
            res += (f_zero[i] * f_one[i - 1]) % mod;
            res += (f_zero[i] * 2 * f_one[i]) % mod;
            res += (f_zero[i] * f_one[i + 1]) % mod;
            res %= mod;
        }
        return res;
    }
};