3739. Count Subarrays With Majority Element II

class BinaryIndexedTree {
    private final int n;
    private final int[] bit;

    public BinaryIndexedTree(int n) {
        this.n = n;
        bit = new int[n + 1];
    }

    public void update(int idx, int val) {
        while (idx <= n) {
            bit[idx] += val;
            idx += idx & -idx;
        }
    }

    public int query(int idx) {
        int sum = 0;
        while (idx > 0) {
            sum += bit[idx];
            idx -= idx & -idx;
        }
        return sum;
    }
}

class Solution {
    public long countMajoritySubarrays(int[] nums, int target) {
        int n = nums.length;

        BinaryIndexedTree bit = new BinaryIndexedTree(2 * n + 1);

        int prefix = n + 1;
        bit.update(prefix, 1);

        long ans = 0;

        for (int x : nums) {
            if (x == target) {
                prefix++;
            } else {
                prefix--;
            }

            ans += bit.query(prefix - 1);
            bit.update(prefix, 1);
        }

        return ans;
    }
}

class BinaryIndexedTree:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)

    def update(self, idx, val):
        while idx <= self.n:
            self.bit[idx] += val
            idx += idx & -idx

    def query(self, idx):
        s = 0
        while idx > 0:
            s += self.bit[idx]
            idx -= idx & -idx
        return s


class Solution:
    def countMajoritySubarrays(self, nums: List[int], target: int) -> int:
        n = len(nums)

        bit = BinaryIndexedTree(2 * n + 1)

        prefix = n + 1
        bit.update(prefix, 1)

        ans = 0

        for x in nums:
            if x == target:
                prefix += 1
            else:
                prefix -= 1

            ans += bit.query(prefix - 1)
            bit.update(prefix, 1)

        return ans

class BinaryIndexedTree {
    int n;
    vector<int> bit;

public:
    BinaryIndexedTree(int n) : n(n), bit(n + 1, 0) {}

    void update(int idx, int val) {
        while (idx <= n) {
            bit[idx] += val;
            idx += idx & -idx;
        }
    }

    int query(int idx) {
        int sum = 0;
        while (idx > 0) {
            sum += bit[idx];
            idx -= idx & -idx;
        }
        return sum;
    }
};

class Solution {
public:
    long long countMajoritySubarrays(vector<int>& nums, int target) {
        int n = nums.size();

        BinaryIndexedTree bit(2 * n + 1);

        int prefix = n + 1;
        bit.update(prefix, 1);

        long long ans = 0;

        for (int x : nums) {
            if (x == target)
                prefix++;
            else
                prefix--;

            ans += bit.query(prefix - 1);
            bit.update(prefix, 1);
        }

        return ans;
    }
};

class BinaryIndexedTree {
    constructor(n) {
        this.n = n;
        this.bit = new Array(n + 1).fill(0);
    }

    update(idx, val) {
        while (idx <= this.n) {
            this.bit[idx] += val;
            idx += idx & -idx;
        }
    }

    query(idx) {
        let sum = 0;
        while (idx > 0) {
            sum += this.bit[idx];
            idx -= idx & -idx;
        }
        return sum;
    }
}

/**
 * @param {number[]} nums
 * @param {number} target
 * @return {number}
 */
var countMajoritySubarrays = function(nums, target) {
    const n = nums.length;

    const bit = new BinaryIndexedTree(2 * n + 1);

    let prefix = n + 1;
    bit.update(prefix, 1);

    let ans = 0;

    for (const x of nums) {
        if (x === target)
            prefix++;
        else
            prefix--;

        ans += bit.query(prefix - 1);
        bit.update(prefix, 1);
    }

    return ans;
};