pythonarrayssliding-windowprefix-sum

Apply Operations to Make All Array Elements Equal to Zero


So I attempted this leetcode problem and I was able to come up with a solution that passes 1017/1026 test cases. The remaining that failed did so due to the time limit exceeding and not incorrectness of the solution. So I am wondering if anyone has any ideas on how to optimize my approach. I know it has to do with the fact that I am finding the min value in the sublist every single time. I'm pretty sure this can be optimized with some prefix sum logic, but I am not sure how to introduce that into my approach.

Here's the question:

You are given a 0-indexed integer array nums and a positive integer k.

You can apply the following operation on the array any number of times:

Choose any subarray of size k from the array and decrease all its elements by 1. Return true if you can make all the array elements equal to 0, or false otherwise.

A subarray is a contiguous non-empty part of an array.

Example 1:

Input: nums = [2,2,3,1,1,0], k = 3 Output: true

Explanation: We can do the following operations:

  • Choose the subarray [2,2,3]. The resulting array will be nums = [1,1,2,1,1,0].
  • Choose the subarray [2,1,1]. The resulting array will be nums = [1,1,1,0,0,0].
  • Choose the subarray [1,1,1]. The resulting array will be nums = [0,0,0,0,0,0].

Example 2:

Input: nums = [1,3,1,1], k = 2 Output: false

Explanation: It is not possible to make all the array elements equal to 0.

Here's my code:

def checkArray(self, nums: List[int], k: int) -> bool:
        i, j = 0, k        
        while i + k <= len(nums):            
            sublist = nums[i:j]  # Make a copy of the sublist            
            smallest = min(sublist)
            for x in range(len(sublist)):                
                sublist[x] -= smallest  # Modify the values in the sublist
                if x > 0 and sublist[x] < sublist[x-1]:
                    return False
            nums[i:j] = sublist  # Assign the modified sublist back to the original list            
            i += 1
            j += 1
        return sum(nums) == 0

And here's my intuition so my approach can be followed:

What I did was use a sliding window. Within each window, we are doing 2 things, we are simulating applying operations that would reduce the elements to 0 and we are also optimising by checking if the window is valid. A window is valid if no element to the left of the currently iterated over element is greater after the simulation is complete. It makes sense because if an element to the left is greater than it, that means that element to the left cannot be reduced to 0 within a window of size k. By simulation, I simply mean instead of continuously subtracting 1 from all elements in the window until the smallest reaches 0, we simply find the smallest element in the window and subtract it from all elements in the window. It yields the same result but more efficiently.


Solution

  • We can take the following greedy approach: whenever we encounter a non-zero element when iterating forwards, we decrease the entire range of size k starting with that element by the value of that element. There is no need to find the minimum value in the range; if there is any value in the range smaller than the value at the left endpoint, it is impossible to make all elements equal to 0. To quickly perform the range update, we can make use of a difference array.

    Sample implementation:

    class Solution:
        def checkArray(self, nums: List[int], k: int) -> bool:
            diff = [0] * (len(nums) + 1)
            for i, v in enumerate(nums):
                if i: diff[i] += diff[i - 1]
                v += diff[i]
                if v < 0: return False
                if v:
                    if i + k > len(nums): return False
                    diff[i] -= v
                    diff[i + k] += v
            return True