Problem Statement
leetcode problem link
Brute Force [TLE]
class Solution:
def countValidSelections(self, nums: List[int]) -> int:
N = len(nums)
res = 0
def dp(curr, is_left):
left = right = 0
if curr < 0 or curr >= N:
return all(x == 0 for x in nums)
elif nums[curr] == 0:
if is_left:
left = dp(curr - 1, is_left)
else:
right = dp(curr + 1, is_left)
elif nums[curr] > 0:
nums[curr] -= 1
if is_left:
right = dp(curr + 1, False)
else:
left = dp(curr - 1, True)
nums[curr] += 1
return left + right
for i in range(N):
if nums[i] == 0:
res += dp(i, True)
res += dp(i, False)
return res
Editorial
Approach 1: Simulation
class Solution:
def countValidSelections(self, nums):
count = 0
nonZeros = sum(1 for x in nums if x > 0)
n = len(nums)
for i in range(n):
if nums[i] == 0:
if self.isValid(nums, nonZeros, i, -1):
count += 1
if self.isValid(nums, nonZeros, i, 1):
count += 1
return count
def isValid(self, nums, nonZeros, start, direction):
temp = nums[:]
curr = start
while nonZeros > 0 and 0 <= curr < len(nums):
if temp[curr] > 0:
temp[curr] -= 1
direction *= -1
if temp[curr] == 0:
nonZeros -= 1
curr += direction
return nonZeros == 0
Approach 2: Prefix Sum
class Solution:
def countValidSelections(self, nums):
n = len(nums)
ans = 0
s = sum(nums)
left, right = 0, s
for i in range(n):
if nums[i] == 0:
if 0 <= left - right <= 1:
ans += 1
if 0 <= right - left <= 1:
ans += 1
else:
left += nums[i]
right -= nums[i]
return ans