LeetCode-Solutions/Algorithms/src/3721. Longest Balanced Subarray II/Code/solution.cpp at main · withaarzoo/LeetCode-Solutions · 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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
/* C++ (optimized, commented) */
#include <bits/stdc++.h>
using namespace std;

struct SegTree
{
    int n;
    vector<int> mn, mx, lazy;
    SegTree(int _n) : n(_n), mn(4 * n, 0), mx(4 * n, 0), lazy(4 * n, 0) {}
    void apply(int idx, int v)
    {
        mn[idx] += v;
        mx[idx] += v;
        lazy[idx] += v;
    }
    void push(int idx)
    {
        if (lazy[idx] != 0)
        {
            apply(idx << 1, lazy[idx]);
            apply(idx << 1 | 1, lazy[idx]);
            lazy[idx] = 0;
        }
    }
    void pull(int idx)
    {
        mn[idx] = min(mn[idx << 1], mn[idx << 1 | 1]);
        mx[idx] = max(mx[idx << 1], mx[idx << 1 | 1]);
    }
    void add_range(int idx, int l, int r, int ql, int qr, int val)
    {
        if (ql > qr)
            return;
        if (ql <= l && r <= qr)
        {
            apply(idx, val);
            return;
        }
        push(idx);
        int mid = (l + r) >> 1;
        if (ql <= mid)
            add_range(idx << 1, l, mid, ql, min(qr, mid), val);
        if (qr > mid)
            add_range(idx << 1 | 1, mid + 1, r, max(ql, mid + 1), qr, val);
        pull(idx);
    }
    // public wrapper
    void add_range(int l, int r, int val)
    {
        if (l > r)
            return;
        add_range(1, 0, n - 1, l, r, val);
    }
    // find rightmost index in [ql, qr] with value == 0, or -1 if none
    int find_rightmost_zero(int idx, int l, int r, int ql, int qr)
    {
        if (ql > qr || qr < l || ql > r)
            return -1;
        if (mn[idx] > 0 || mx[idx] < 0)
            return -1; // no zero inside
        if (l == r)
        {
            if (mn[idx] == 0)
                return l;
            return -1;
        }
        push(idx);
        int mid = (l + r) >> 1;
        // try right child first to get rightmost
        if (qr > mid)
        {
            int res = find_rightmost_zero(idx << 1 | 1, mid + 1, r, max(ql, mid + 1), qr);
            if (res != -1)
                return res;
        }
        if (ql <= mid)
        {
            return find_rightmost_zero(idx << 1, l, mid, ql, min(qr, mid));
        }
        return -1;
    }
    int find_rightmost_zero(int ql, int qr)
    {
        if (ql > qr)
            return -1;
        return find_rightmost_zero(1, 0, n - 1, ql, qr);
    }
};

class Solution
{
public:
    int longestBalanced(vector<int> &nums)
    {
        int n = nums.size();
        unordered_map<int, vector<int>> pos;
        pos.reserve(n * 2);
        for (int i = 0; i < n; ++i)
            pos[nums[i]].push_back(i);

        SegTree st(n);
        // initial: for each value, add sign to [firstPos, n-1]
        for (auto &kv : pos)
        {
            int val = kv.first;
            int sign = (val & 1) ? 1 : -1;
            int p = kv.second[0];
            st.add_range(p, n - 1, sign);
        }

        // pointers to current first occurrence for each value
        unordered_map<int, int> ptr;
        ptr.reserve(pos.size() * 2);
        for (auto &kv : pos)
            ptr[kv.first] = 0;

        int ans = 0;
        for (int l = 0; l < n; ++l)
        {
            int r = st.find_rightmost_zero(l, n - 1);
            if (r != -1)
                ans = max(ans, r - l + 1);

            int x = nums[l];
            int pIndex = ptr[x]; // should point to l
            // move pointer forward
            ptr[x] = pIndex + 1;
            int nextPos = (ptr[x] < (int)pos[x].size()) ? pos[x][ptr[x]] : n;
            int sign = (x & 1) ? 1 : -1;
            // net effect: apply -sign to range [l, nextPos-1]
            int L = l, R = nextPos - 1;
            if (L <= R)
                st.add_range(L, R, -sign);
        }
        return ans;
    }
};