LeetCode：操作使数组全为零的解法优化求助

问题描述 #

给定0索引整数数组nums和正整数k，可执行任意次数操作：选择大小为k的子数组，将其中所有元素减1。判断能否使所有元素变为0。

示例1：
输入：nums = [2,2,3,1,1,0], k = 3
输出：true
解释：执行三次操作可将数组变为全零。

示例2：
输入：nums = [1,3,1,1], k = 2
输出：false
解释：无法将数组变为全零。

现有解法及问题 #

当前滑动窗口解法通过了1017/1026个测试用例，但剩余用例超时。核心问题是每次窗口都要计算子数组最小值，且直接修改数组的操作效率太低。代码如下：

def checkArray(self, nums: List[int], k: int) -> bool:
        i, j = 0, k        
        while i + k <= len(nums):            
            sublist = nums[i:j]  # 复制子数组            
            smallest = min(sublist)
            for x in range(len(sublist)):                
                sublist[x] -= smallest  # 修改子数组元素
                if x > 0 and sublist[x] < sublist[x-1]:
                    return False
            nums[i:j] = sublist  # 将修改后的子数组写回原数组            
            i += 1
            j += 1
        return sum(nums) == 0

思路说明：
采用滑动窗口思路，每个窗口内做两件事：

• 模拟将元素减至0的操作（直接减去窗口最小值而非多次减1）
• 验证窗口有效性（模拟后左侧元素不大于当前元素，否则左侧元素无法在k大小窗口内减至0）

前缀和优化思路 #

可以用差分+前缀和的方式替代原逻辑，把时间复杂度降到O(n)，彻底解决超时问题。

核心逻辑 #

无需真的修改每个窗口的元素，而是用差分数组记录操作的影响范围：

1. 遍历数组时，用前缀和维护当前位置累计被操作的次数current_ops，计算当前元素的实际值nums[i] - current_ops。
2. 对实际值的处理：
   • 若实际值小于0：直接返回False，说明之前操作减多了，无法恢复。
   • 若实际值等于0：无需操作，继续遍历。
   • 若实际值大于0：需要执行actual_val次操作，覆盖i到i+k-1的子数组。此时在差分数组diff[i]加actual_val，在diff[i+k]减actual_val（若i+k不超过数组长度），以此标记操作的影响范围。
3. 遍历到数组最后k-1个元素时，无法再发起新操作（窗口会越界），若此时实际值不为0，直接返回False。
4. 遍历结束后返回True。

优化后的代码 #

def checkArray(self, nums: List[int], k: int) -> bool:
    n = len(nums)
    diff = [0] * (n + 1)  # 差分数组长度设为n+1，避免越界
    current_ops = 0
    
    for i in range(n):
        current_ops += diff[i]
        actual_val = nums[i] - current_ops
        
        if actual_val < 0:
            return False
        if actual_val == 0:
            continue
        
        # 剩余元素不足k个，无法发起操作
        if i + k > n:
            return False
        # 标记操作的影响范围
        diff[i] += actual_val
        diff[i + k] -= actual_val
        current_ops += actual_val  # 同步更新当前累计操作次数
    
    return True

效率说明 #

• 避免了原解法中每次窗口的最小值计算、数组拷贝和修改操作，所有步骤都是O(1)级别的遍历。
• 差分数组的前缀和天然维护了每个位置的累计操作次数，完美替代滑动窗口中重复修改元素的逻辑。
• 时间复杂度O(n)，空间复杂度O(n)（也可优化为O(1)，用变量记录后续需要抵消的操作次数，O(n)实现更直观）。

内容的提问来源于stack exchange，提问作者Duck Dodgers