forked from doocs/leetcode
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathSolution.py
24 lines (23 loc) · 850 Bytes
/
Solution.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
def numberOfStableArrays(self, zero: int, one: int, limit: int) -> int:
@cache
def dfs(i: int, j: int, k: int) -> int:
if i == 0:
return int(k == 1 and j <= limit)
if j == 0:
return int(k == 0 and i <= limit)
if k == 0:
return (
dfs(i - 1, j, 0)
+ dfs(i - 1, j, 1)
- (0 if i - limit - 1 < 0 else dfs(i - limit - 1, j, 1))
)
return (
dfs(i, j - 1, 0)
+ dfs(i, j - 1, 1)
- (0 if j - limit - 1 < 0 else dfs(i, j - limit - 1, 0))
)
mod = 10**9 + 7
ans = (dfs(zero, one, 0) + dfs(zero, one, 1)) % mod
dfs.cache_clear()
return ans