I use dynamic programming where dp[r] is the number of ways to partition the first r elements (r from 0..n) and dp[0]=1. For each right end r (1-based for prefix length) I find the smallest l (0-based index) such that the subarray nums[l..r-1] has max - min ≤ k; then dp[r] = sum(dp[i] for i in [l..r-1]). I compute that range-sum in O(1) with a running prefix-sum array pref. To get every l in amortized O(1) time I maintain a sliding window with two monotonic deques for max and min and advance the left pointer until the window becomes valid.
class Solution:
def countPartitions(self, nums: List[int], k: int) -> int:
MOD = 10**9 + 7
n = len(nums)
dp0 = 1
pref_prev = dp0
left = 0
maxdq = deque()
mindq = deque()
dp_r = 0
for r in range(1, n + 1):
x = nums[r - 1]
while maxdq and maxdq[-1][0] < x:
maxdq.pop()
maxdq.append((x, r - 1))
while mindq and mindq[-1][0] > x:
mindq.pop()
mindq.append((x, r - 1))
while maxdq and mindq and (maxdq[0][0] - mindq[0][0] > k):
if maxdq and maxdq[0][1] == left:
maxdq.popleft()
if mindq and mindq[0][1] == left:
mindq.popleft()
left += 1
if left == 0:
dp_r = pref_prev
else:
pass
break
pref = [0] * (n + 1)
pref[0] = 1
left = 0
maxdq.clear()
mindq.clear()
for r in range(1, n + 1):
x = nums[r - 1]
while maxdq and maxdq[-1][0] < x:
maxdq.pop()
maxdq.append((x, r - 1))
while mindq and mindq[-1][0] > x:
mindq.pop()
mindq.append((x, r - 1))
while maxdq and mindq and (maxdq[0][0] - mindq[0][0] > k):
if maxdq and maxdq[0][1] == left:
maxdq.popleft()
if mindq and mindq[0][1] == left:
mindq.popleft()
left += 1
dp_r = (pref[r - 1] - (pref[left - 1] if left > 0 else 0)) % MOD
pref[r] = (pref[r - 1] + dp_r) % MOD
return dp_r % MOD- Time: O(n)
- Space: O(n)
