Make Array Elements Equal to Zero

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.

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.

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();
    }
}