Skip to content

3129. Find All Possible Stable Binary Arrays I

  • Time: $O(mn)$
  • Space: $O(mn)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution {
 public:
  // Same as 3129. Find All Possible Stable Binary Arrays I
  int numberOfStableArrays(int zero, int one, int limit) {
    constexpr int kMod = 1'000'000'007;
    // dp[i][j][k] := the number of stable arrays, where the number of
    // ocurrences of 0 is i and the number of ocurrences of 1 is j and the last
    // number is k (0/1)
    vector<vector<vector<long>>> dp(
        zero + 1, vector<vector<long>>(one + 1, vector<long>(2)));

    for (int i = 0; i <= min(zero, limit); ++i)
      dp[i][0][0] = 1;

    for (int j = 0; j <= min(one, limit); ++j)
      dp[0][j][1] = 1;

    for (int i = 1; i <= zero; ++i)
      for (int j = 1; j <= one; ++j) {
        dp[i][j][0] = (dp[i - 1][j][0] + dp[i - 1][j][1] -
                       (i - limit < 1 ? 0 : dp[i - limit - 1][j][1]) + kMod) %
                      kMod;
        dp[i][j][1] = (dp[i][j - 1][0] + dp[i][j - 1][1] -
                       (j - limit < 1 ? 0 : dp[i][j - limit - 1][0]) + kMod) %
                      kMod;
      }

    return (dp[zero][one][0] + dp[zero][one][1]) % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
  // Same as 3129. Find All Possible Stable Binary Arrays I
  public int numberOfStableArrays(int zero, int one, int limit) {
    final int MOD = 1_000_000_007;
    // dp[i][j][k] := the number of stable arrays, where the number of
    // occurrences of 0 is i and the number of occurrences of 1 is j and the last
    // number is k (0/1)
    long[][][] dp = new long[zero + 1][one + 1][2];

    for (int i = 0; i <= Math.min(zero, limit); ++i)
      dp[i][0][0] = 1;

    for (int j = 0; j <= Math.min(one, limit); ++j)
      dp[0][j][1] = 1;

    for (int i = 1; i <= zero; ++i)
      for (int j = 1; j <= one; ++j) {
        dp[i][j][0] = (dp[i - 1][j][0] + dp[i - 1][j][1] -
                       (i - limit < 1 ? 0 : dp[i - limit - 1][j][1]) + MOD) %
                      MOD;
        dp[i][j][1] = (dp[i][j - 1][0] + dp[i][j - 1][1] -
                       (j - limit < 1 ? 0 : dp[i][j - limit - 1][0]) + MOD) %
                      MOD;
      }

    return (int) ((dp[zero][one][0] + dp[zero][one][1]) % MOD);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution:
  # Same as 3129. Find All Possible Stable Binary Arrays I
  def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
    MOD = 1_000_000_007
    # dp[i][j][k] := the number of stable arrays, where the number of
    # occurrences of 0 is i and the number of occurrences of 1 is j and the last
    # number is k (0/1)
    dp = [[[0] * 2
           for _ in range(one + 1)]
          for _ in range(zero + 1)]

    for i in range(min(zero, limit) + 1):
      dp[i][0][0] = 1

    for j in range(min(one, limit) + 1):
      dp[0][j][1] = 1

    for i in range(1, zero + 1):
      for j in range(1, one + 1):
        dp[i][j][0] = (
            dp[i - 1][j][0] + dp[i - 1][j][1] -
            (dp[i - limit - 1][j][1] if i - limit >= 1 else 0) + MOD) % MOD
        dp[i][j][1] = (
            dp[i][j - 1][0] + dp[i][j - 1][1] -
            (dp[i][j - limit - 1][0] if j - limit >= 1 else 0) + MOD) % MOD

    return (dp[zero][one][0] + dp[zero][one][1]) % MOD