Find All Possible Stable Binary Arrays II

Find All Possible Stable Binary Arrays II

You are given 3 positive integers zero, one, and limit.

A binary array arr is called stable if:

• The number of occurrences of 0 in arr is exactly zero.
• The number of occurrences of 1 in arr is exactly one.
• Each subarray of arr with a size greater than limit must contain both 0 and 1.

Return the total number of stable binary arrays.

Since the answer may be very large, return it modulo 109 + 7.

 

Example 1:

Input: zero = 1, one = 1, limit = 2

Output: 2

Explanation:

The two possible stable binary arrays are [1,0] and [0,1].

Example 2:

Input: zero = 1, one = 2, limit = 1

Output: 1

Explanation:

The only possible stable binary array is [1,0,1].

Example 3:

Input: zero = 3, one = 3, limit = 2

Output: 14

Explanation:

All the possible stable binary arrays are [0,0,1,0,1,1], [0,0,1,1,0,1], [0,1,0,0,1,1], [0,1,0,1,0,1], [0,1,0,1,1,0], [0,1,1,0,0,1], [0,1,1,0,1,0], [1,0,0,1,0,1], [1,0,0,1,1,0], [1,0,1,0,0,1], [1,0,1,0,1,0], [1,0,1,1,0,0], [1,1,0,0,1,0], and [1,1,0,1,0,0].

 

Constraints:

• 1 <= zero, one, limit <= 1000

 

解题思路

  一开始状态定义的不好，导致优化起来巨复杂。为了方便这里重新定义符号 $a = \text{zero}$，$b = \text{one}$，$k = \text{limit}$，$n = a + b$。

  容易知道当 $1$ 的位置确定后，$0$ 的位置也就确定了，所以我们可以只关注 $1$ 的位置。定义 $f(i,j)$ 表示在前 $j$ 个位置中放入了 $i$ 个 $1$ 且第 $j$ 个位置是 $1$ 的方案数。根据最后一段的 $1$ 的个数及其前一段 $0$ 的个数进行状态划分。

  具体来说如果最后一段有 $u$ 个 $1$，显然 $1 \leq u \leq k$，其前一段 $0$ 的个数为 $v$，同样满足 $1 \leq v \leq k$，则有 $f(i, j) \gets f(i - u, j - u - v)$。所以状态转移方程就是： $$f(i,j) = \sum_{u = 1}^{k}{\sum_{v = 1}^{k}{f(i-u, j-u-v)}}$$

  显然直接暴力 dp 的话时间复杂度为 $O(b \, n \, k^2)$。我们把上面的式子展开：

$$\begin{align*} f(i-1,j-2)+f(i-1,j-3)+f(i-1,j-4)+\cdots&+f(i-1,j-1-k) \\ +f(i-2,j-3)+f(i-2,j-4)+\cdots&+f(i-2,j-1-k)+f(i-2,j-2-k) \\ +f(i-3,j-4)+\cdots&+f(i-3,j-1+k)+f(i-3,j-2-k)+f(i-3,j-3-k) \\ &\vdots \\ &+f(i-k,j-1-k)+f(i-k,j-2-k)+f(i-k,j-3-k)+\cdots+f(i-k,j-k-k) \\ \end{align*}$$

  可以发现有点类似于平行四边形。再考虑 $f(i,j+1)$：

$$\begin{align*} {\color{Red} {f(i-1,j-1)}}+f(i-1,j-2)+f(i-1,j-3)+\cdots&+f(i-1,j-k) \\ +{\color{Red}{f(i-2,j-2)}}+f(i-2,j-3)+\cdots&+f(i-2,j-k)+f(i-2,j-1-k) \\ +{\color{Red}{f(i-3,j-3)}} +\cdots&+f(i-3,j+k)+f(i-3,j-1-k)+f(i-3,j-2-k) \\ \vdots \\ &+{\color{Red}{f(i-k,j-k)}}+f(i-k,j-1-k)+f(i-k,j-2-k)+\cdots+f(i-k,j+1-k-k) \end{align*}$$

  与 $f(i,j)$ 相比 $f(i,j+1)$ 只多出了红色的部分，可以看作是以 $f(i-1,j-1)$ 为端点的对角线上 $k$ 个 $f$ 的和，即 $\sum\limits_{u=0}^{k-1}{f(i-1-u,j-1-u)}$。另外减少的部分是 $\sum\limits_{u=0}^{k-1}{f(i-1-u,j-1-k-u)}$。

  为此我们可以维护一个 $g(i,j) = \sum\limits_{u=0}^{k-1}{f(i-u,j-u)}$，转移方程为 $g(i,j) = g(i-1,j-1) + f(i,j) - f(i-k,j-k)$。

  这样当我们确定 $i$ 再枚举 $j$ 时，可以维护一个大小为 $k$ 的滑动窗口，并用 $s$ 来维护窗口内的 $g$ 的和。当枚举到 $i$ 和 $j$，有 $s \gets s + g(i-1,j-2) - g(i-1,j-k-2)$，然后有 $f(i,j) = s$。

  另外边界情况，即枚举到 $i$ 时前 $i$ 个位置均为 $1$，此时有 $f(i,i) = 1$，当然只有 $i \leq k$ 时才成立。剩下的细节见代码。

  最后根据最后一段 $0$ 的不同个数来统计答案，即 $\sum\limits_{i=0}^{\min\{k,n\}}{f(b,n-i)}$。

  AC 代码如下，时间复杂度为 $O(b \times n)$：

class Solution {
public:
    int mod = 1e9 + 7;
    
    int numberOfStableArrays(int a, int b, int k) {
        int n = a + b;
        vector<vector<int>> f(b + 1, vector<int>(n + 1)), g(b + 1, vector<int>(n + 1));
        f[0][0] = g[0][0] = 1;
        for (int i = 1; i <= b; i++) {
            if (i <= k) f[i][i] = 1;
            for (int j = i + 1, s = 0; j <= n; j++) {
                s = (s + g[i - 1][j - 2]) % mod;
                if (j - k - 2 >= 0) s = (s - g[i - 1][j - k - 2]) % mod;
                f[i][j] = s;
            }
            // g[i][0] = f[i][0];
            for (int j = 1; j <= n; j++) {
                g[i][j] = (g[i - 1][j - 1] + f[i][j]) % mod;
                if (i - k >= 0 && j - k >= 0) g[i][j] = (g[i][j] - f[i - k][j - k]) % mod;
            }
        }
        int ret = 0;
        for (int i = 0; i <= k && i <= n; i++) {
            ret = (ret + f[b][n - i]) % mod;
        }
        ret = (ret + mod) % mod;
        return ret;
    }
};

   事实上如果把状态定义成 $f(i,j,0/1)$ 表示放置了 $i$ 个 $0$，$j$ 个 $1$ 且最后一段是 $0/1$ 的方案数，那么就会变得十分简单了。反正我是想不到的，关键点应该是要发现对于任意一个合法方案，每一段的 $0$ 和 $1$ 是交替出现的。

  状态转移方程就很简单了。

$$\begin{cases} f(i,j,0) = \sum\limits_{u=1}^{k}{f(i-u,j,1)} \\ f(i,j,1) = \sum\limits_{u=1}^{k}{f(i,j-u,0)} \\ \end{cases}$$

  优化的地方就很容易看到了，也是用一个大小为 $k$ 的滑动窗口来维护 $f$。其中对于 $f(i,j,0)$，应该在每一列开一个单调队列。而对于 $f(i,j,1)$，则应该在每一行开一个单调队列。

  最后答案就是 $f(a,b,0) + f(a,b,1)$。

  AC 代码如下，时间复杂度为 $O(a \times b)$：

class Solution {
public:
    int mod = 1e9 + 7;
    
    int numberOfStableArrays(int a, int b, int k) {
        vector<vector<vector<int>>> f(a + 1, vector<vector<int>>(b + 1, vector<int>(2)));
        vector<int> c(b + 1), r(a + 1);
        for (int i = 1; i <= k && i <= a; i++) {
            f[i][0][0] = 1;
        }
        for (int i = 1; i <= k && i <= b; i++) {
            f[0][i][1] = 1;
        }
        for (int i = 1; i <= a; i++) {
            for (int j = 1; j <= b; j++) {
                if (i - 1 >= 0) c[j] = (c[j] + f[i - 1][j][1]) % mod;
                if (i - k - 1 >= 0) c[j] = (c[j] - f[i - k - 1][j][1]) % mod;
                f[i][j][0] = c[j];
                if (j - 1 >= 0) r[i] = (r[i] + f[i][j - 1][0]) % mod;
                if (j - k - 1 >= 0) r[i] = (r[i] - f[i][j - k - 1][0]) % mod;
                f[i][j][1] = r[i];
            }
        }
        return ((f[a][b][0] + f[a][b][1]) % mod + mod) % mod;
    }
};

 

参考资料

  第 129 场力扣夜喵双周赛 - 力扣（LeetCode）：https://leetcode.cn/circle/discuss/SZyTo4/view/mv9LwV/