LeetCode/3354-make-array-elements-equal-to-zero.cpp at master · benwu25/LeetCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
#include <vector>
#include <set>
using namespace std;

class Solution {
public:
    bool simulate(vector<int> nums, int curr, bool left) {
        while (0 <= curr && curr < nums.size()) {
            if (nums[curr] == 0 && left) {
                --curr;
            } else if (nums[curr] == 0 && !left) {
                ++curr;
            } else {
                --nums[curr];
                left = !left;
                if (left)
                    --curr;
                else
                    ++curr;
            }
        }
        for (auto &elem : nums)
            if (elem != 0)
                return false;
        return true;
    }

    int countValidSelections(vector<int>& nums) {
        set<int> zeros;
        for (int i = 0; i < nums.size(); ++i) {
            if (nums[i] == 0)
                zeros.insert(i);
        }

        int valid = 0;
        for (auto &selection_start : zeros) {
            if (simulate(nums, selection_start, true))
                ++valid;
            if (simulate(nums, selection_start, false))
                ++valid;
        }
        return valid;
    }
};

/*

    find all indices where nums[i] == 0
    for each index,

*/