Scale Engineer

Make Array Elements Equal to Zero

Easy
#2980Time: O(Z * (n log n + S * log n)), where Z is the number of zeros, S is the sum of elements, and n is the length. For each of the `Z` starting configurations, we spend `O(n log n)` to initialize the `TreeSet`. The simulation involves `S` decrements, and each step (decrement or jump) takes `O(log n)` time. This complexity is efficient enough for the given constraints.Space: O(n), where n is the length of `nums`. Space is used for the array copy and the `TreeSet`, which can store up to `n` indices.
Patterns
Prefix Sum
Data structures
Array

Prompt

You are given an integer array nums.

Start by selecting a starting position curr such that nums[curr] == 0, and choose a movement direction of either left or right.

After that, you repeat the following process:

  • If curr is out of the range [0, n - 1], this process ends.
  • If nums[curr] == 0, move in the current direction by incrementing curr if you are moving right, or decrementing curr if you are moving left.
  • Else if nums[curr] > 0:
    • Decrement nums[curr] by 1.
    • Reverse your movement direction (left becomes right and vice versa).
    • Take a step in your new direction.

A selection of the initial position curr and movement direction is considered valid if every element in nums becomes 0 by the end of the process.

Return the number of possible valid selections.

 

Example 1:

Input: nums = [1,0,2,0,3]

Output: 2

Explanation:

The only possible valid selections are the following:

  • Choose curr = 3, and a movement direction to the left.
    • [1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,1,0,3] -> [1,0,1,0,3] -> [1,0,1,0,2] -> [1,0,1,0,2] -> [1,0,0,0,2] -> [1,0,0,0,2] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,0].
  • Choose curr = 3, and a movement direction to the right.
    • [1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,2,0,2] -> [1,0,2,0,2] -> [1,0,1,0,2] -> [1,0,1,0,2] -> [1,0,1,0,1] -> [1,0,1,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [0,0,0,0,0].

Example 2:

Input: nums = [2,3,4,0,4,1,0]

Output: 0

Explanation:

There are no possible valid selections.

 

Constraints:

  • 1 <= nums.length <= 100
  • 0 <= nums[i] <= 100
  • There is at least one element i where nums[i] == 0.

Approaches

2 approaches with complexity analysis and trade-offs.

The brute-force simulation can be slow because of the time spent traversing blocks of zeros one step at a time. We can optimize this by skipping over consecutive zeros in a single operation. A TreeSet is an effective data structure for this purpose, as it can maintain the sorted indices of non-zero elements and find the next one in a given direction in logarithmic time.

Algorithm

  1. The overall structure is the same as the brute-force approach: iterate through all (start_position, start_direction) pairs.
  2. For each simulation, instead of a simple array, use a TreeSet to store the indices of all non-zero elements. This allows for efficient queries of the next non-zero index.
  3. Initialize the simulation with a copy of nums, the TreeSet of non-zero indices, curr, and dir.
  4. Loop as long as curr is within bounds: a. If nums_copy[curr] > 0: i. Decrement nums_copy[curr]. ii. If nums_copy[curr] becomes 0, remove curr from the TreeSet. iii. Reverse direction and move: dir *= -1, curr += dir. b. If nums_copy[curr] == 0: i. Instead of moving one step at a time, find the next non-zero element in the current direction using the TreeSet. ii. If moving right (dir=1), find the smallest index in the TreeSet that is greater than curr. This is done via treeSet.ceiling(curr + 1). iii. If moving left (dir=-1), find the largest index in the TreeSet that is less than curr. This is done via treeSet.floor(curr - 1). iv. If a next non-zero index is found, jump curr to that index. Otherwise, jump curr out of bounds.
  5. After the loop terminates, the selection is valid if and only if the TreeSet of non-zero indices is empty.

Walkthrough

In this approach, we enhance the simulation by maintaining a TreeSet of indices where the array value is positive. When the simulation is at a position curr with a value of 0, instead of moving one step, we query the TreeSet to find the index of the next non-zero element in the current direction. This allows us to jump over entire blocks of zeros instantly.

For a rightward movement, treeSet.ceiling(curr + 1) gives us the next non-zero index. For a leftward movement, treeSet.floor(curr - 1) provides the previous non-zero index. If no such index exists, it means we've moved past all remaining non-zero numbers in that direction, and we can jump curr out of bounds to terminate that path segment.

When a non-zero element at an index is decremented to zero, we must remove that index from the TreeSet. This keeps our set of non-zero indices up-to-date. This optimization significantly reduces the number of steps in the simulation, especially for sparse arrays with large blocks of zeros.

class Solution {    public int makeArrayElementsEqualtoZero(int[] nums) {        int n = nums.length;        java.util.List<Integer> zeroIndices = new java.util.ArrayList<>();        for (int i = 0; i < n; i++) {            if (nums[i] == 0) {                zeroIndices.add(i);            }        }         int validCount = 0;        for (int startIdx : zeroIndices) {            if (isSelectionValid(nums.clone(), startIdx, 1)) {                validCount++;            }            if (isSelectionValid(nums.clone(), startIdx, -1)) {                validCount++;            }        }        return validCount;    }     private boolean isSelectionValid(int[] currentNums, int startCurr, int startDir) {        int n = currentNums.length;        java.util.TreeSet<Integer> nonZeroIndices = new java.util.TreeSet<>();        for (int i = 0; i < n; i++) {            if (currentNums[i] > 0) {                nonZeroIndices.add(i);            }        }         if (nonZeroIndices.isEmpty()) {            return true;        }         int curr = startCurr;        int dir = startDir;         while (curr >= 0 && curr < n) {            if (currentNums[curr] > 0) {                currentNums[curr]--;                if (currentNums[curr] == 0) {                    nonZeroIndices.remove(curr);                }                dir *= -1;                curr += dir;            } else { // currentNums[curr] == 0                if (dir == 1) {                    Integer nextNonZero = nonZeroIndices.ceiling(curr + 1);                    curr = (nextNonZero != null) ? nextNonZero : n;                } else { // dir == -1                    Integer prevNonZero = nonZeroIndices.floor(curr - 1);                    curr = (prevNonZero != null) ? prevNonZero : -1;                }            }        }         return nonZeroIndices.isEmpty();    }}

Complexity

Time

O(Z * (n log n + S * log n)), where Z is the number of zeros, S is the sum of elements, and n is the length. For each of the `Z` starting configurations, we spend `O(n log n)` to initialize the `TreeSet`. The simulation involves `S` decrements, and each step (decrement or jump) takes `O(log n)` time. This complexity is efficient enough for the given constraints.

Space

O(n), where n is the length of `nums`. Space is used for the array copy and the `TreeSet`, which can store up to `n` indices.

Trade-offs

Pros

  • Significantly faster than the naive simulation, especially for sparse arrays.

  • The time complexity is well within typical limits for competitive programming platforms.

Cons

  • Slightly more complex to implement due to the use of a TreeSet.

Solutions

class Solution {public  int countValidSelections(int[] nums) {    int s = Arrays.stream(nums).sum();    int ans = 0, l = 0;    for (int x : nums) {      if (x != 0) {        l += x;      } else if (l * 2 == s) {        ans += 2;      } else if (Math.abs(l * 2 - s) <= 1) {        ++ans;      }    }    return ans;  }}

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem