Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code : Count Good Subarrays

inline bool is_bit_set(int x, int i) { return ((1 << i) & x) != 0; }

class Solution {
  public:
    long long countGoodSubarrays(vector<int> &nums);
};

long long Solution ::countGoodSubarrays(vector<int> &a) {
    int n = a.size();

    vector<vector<int>> nearest_one_on_left(n, vector<int>(30, -1));
    vector<vector<int>> nearest_one_on_right(n, vector<int>(30, n));
    vector<int> prev_occ(n, -1);

    map<int, int> last_seen;
    for (int i = 0; i < n; i++) {
        if (last_seen.count(a[i])) {
            prev_occ[i] = last_seen[a[i]];
        }
        last_seen[a[i]] = i;
    }

    // nearest_one_on_left[i][bit] is the index of the nearest one on the left
    // of a[i] in the bit position bit
    for (int bit = 0; bit < 30; bit++) {
        int last = -1;
        for (int i = 0; i < n; i++) {
            nearest_one_on_left[i][bit] = last;
            if (is_bit_set(a[i], bit)) {
                last = i;
            }
        }
    }

    // nearest_one_on_right[i][bit] is the index of the nearest one on the right
    // of a[i] in the bit position bit
    for (int bit = 0; bit < 30; bit++) {
        int last = n;
        for (int i = n - 1; i >= 0; i--) {
            nearest_one_on_right[i][bit] = last;
            if (is_bit_set(a[i], bit)) {
                last = i;
            }
        }
    }

    long long ans = 0;
    for (int i = 0; i < n; i++) {
        int left_boundary = prev_occ[i];
        int right_boundary = n;
        for (int bit = 0; bit < 30; bit++) {
            if (!is_bit_set(a[i], bit)) {
                left_boundary = max(left_boundary, nearest_one_on_left[i][bit]);
                right_boundary =
                    min(right_boundary, nearest_one_on_right[i][bit]);
            }
        }
        ans += 1LL * (i - left_boundary) * (right_boundary - i);
    }

    return ans;
}
class Solution {
  public:
    long long countGoodSubarrays(vector<int> &nums) {
        int sz = nums.size();
        vector<int> left_bound(sz), right_bound(sz);
        int prev_bit_pos[30];
        for (int i = 0; i < 30; ++i)
            prev_bit_pos[i] = -1;
        unordered_map<int, int> last_val_idx;

        for (int i = 0; i < sz; ++i) {
            int l_limit = -1;
            for (int bit = 0; bit < 30; ++bit) {
                if (((nums[i] >> bit) & 1) == 0) {
                    if (prev_bit_pos[bit] > l_limit) {
                        l_limit = prev_bit_pos[bit];
                    }
                }
            }
            if (last_val_idx.count(nums[i])) {
                if (last_val_idx[nums[i]] > l_limit) {
                    l_limit = last_val_idx[nums[i]];
                }
            }
            left_bound[i] = l_limit + 1;

            for (int bit = 0; bit < 30; ++bit) {
                if ((nums[i] >> bit) & 1) {
                    prev_bit_pos[bit] = i;
                }
            }
            last_val_idx[nums[i]] = i;
        }

        int next_bit_pos[30];
        for (int i = 0; i < 30; ++i)
            next_bit_pos[i] = sz;

        for (int i = sz - 1; i >= 0; --i) {
            int r_limit = sz;
            for (int bit = 0; bit < 30; ++bit) {
                if (((nums[i] >> bit) & 1) == 0) {
                    if (next_bit_pos[bit] < r_limit) {
                        r_limit = next_bit_pos[bit];
                    }
                }
            }
            right_bound[i] = r_limit - 1;

            for (int bit = 0; bit < 30; ++bit) {
                if ((nums[i] >> bit) & 1) {
                    next_bit_pos[bit] = i;
                }
            }
        }

        long long total_good = 0;
        for (int i = 0; i < sz; ++i) {
            long long left_options = i - left_bound[i] + 1;
            long long right_options = right_bound[i] - i + 1;
            total_good += left_options * right_options;
        }

        return total_good;
    }
};