README_EN.md

2031. Count Subarrays With More Ones Than Zeros

中文文档

Description

You are given a binary array nums containing only the integers 0 and 1. Return the number of subarrays in nums that have more 1's than 0's. Since the answer may be very large, return it modulo 109 + 7.

A subarray is a contiguous sequence of elements within an array.

 

Example 1:

Input: nums = [0,1,1,0,1]
Output: 9
Explanation:
The subarrays of size 1 that have more ones than zeros are: [1], [1], [1]
The subarrays of size 2 that have more ones than zeros are: [1,1]
The subarrays of size 3 that have more ones than zeros are: [0,1,1], [1,1,0], [1,0,1]
The subarrays of size 4 that have more ones than zeros are: [1,1,0,1]
The subarrays of size 5 that have more ones than zeros are: [0,1,1,0,1]

Example 2:

Input: nums = [0]
Output: 0
Explanation:
No subarrays have more ones than zeros.

Example 3:

Input: nums = [1]
Output: 1
Explanation:
The subarrays of size 1 that have more ones than zeros are: [1]

 

Constraints:

• 1 <= nums.length <= 105
• 0 <= nums[i] <= 1

Solutions

Binary Indexed Tree.

Python3

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

    @staticmethod
    def lowbit(x):
        x += int(1e5 + 1)
        return x & -x

    def update(self, x, delta):
        x += int(1e5 + 1)
        while x <= self.n:
            self.c[x] += delta
            x += BinaryIndexedTree.lowbit(x)

    def query(self, x):
        x += int(1e5 + 1)
        s = 0
        while x > 0:
            s += self.c[x]
            x -= BinaryIndexedTree.lowbit(x)
        return s


class Solution:
    def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int:
        n = len(nums)
        s = [0]
        for v in nums:
            s.append(s[-1] + (v or -1))
        tree = BinaryIndexedTree(n + 1)
        MOD = int(1e9 + 7)
        ans = 0
        for v in s:
            ans = (ans + tree.query(v - 1)) % MOD
            tree.update(v, 1)
        return ans

Java

class BinaryIndexedTree {
    private int n;
    private int[] c;

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

    public void update(int x, int delta) {
        x += (int) 1e5 + 1;
        while (x <= n) {
            c[x] += delta;
            x += lowbit(x);
        }
    }

    public int query(int x) {
        x += (int) 1e5 + 1;
        int s = 0;
        while (x > 0) {
            s += c[x];
            x -= lowbit(x);
        }
        return s;
    }

    public static int lowbit(int x) {
        x += (int) 1e5 + 1;
        return x & -x;
    }
}

class Solution {
    private static final int MOD = (int) 1e9 + 7;

    public int subarraysWithMoreZerosThanOnes(int[] nums) {
        int n = nums.length;
        int[] s = new int[n + 1];
        for (int i = 0; i < n; ++i) {
            s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1);
        }
        BinaryIndexedTree tree = new BinaryIndexedTree(n + 1);
        int ans = 0;
        for (int v : s) {
            ans = (ans + tree.query(v - 1)) % MOD;
            tree.update(v, 1);
        }
        return ans;
    }
}

C++

class BinaryIndexedTree {
public:
    int n;
    vector<int> c;

    BinaryIndexedTree(int _n): n(_n + 1e5 + 1), c(_n + 1 + 1e5 + 1){}

    void update(int x, int delta) {
        x += 1e5 + 1;
        while (x <= n)
        {
            c[x] += delta;
            x += lowbit(x);
        }
    }

    int query(int x) {
        x += 1e5 + 1;
        int s = 0;
        while (x > 0)
        {
            s += c[x];
            x -= lowbit(x);
        }
        return s;
    }

    int lowbit(int x) {
        x += 1e5 + 1;
        return x & -x;
    }
};

class Solution {
public:
    int subarraysWithMoreZerosThanOnes(vector<int>& nums) {
        int n = nums.size();
        vector<int> s(n + 1);
        for (int i = 0; i < n; ++i) s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1);
        BinaryIndexedTree* tree = new BinaryIndexedTree(n + 1);
        int ans = 0;
        const int MOD = 1e9 + 7;
        for (int v : s)
        {
            ans = (ans + tree->query(v - 1)) % MOD;
            tree->update(v, 1);
        }
        return ans;
    }
};

Go

type BinaryIndexedTree struct {
	n int
	c []int
}

func newBinaryIndexedTree(n int) *BinaryIndexedTree {
	n += 1e5 + 1
	c := make([]int, n+1)
	return &BinaryIndexedTree{n, c}
}

func (this *BinaryIndexedTree) lowbit(x int) int {
	x += 1e5 + 1
	return x & -x
}

func (this *BinaryIndexedTree) update(x, delta int) {
	x += 1e5 + 1
	for x <= this.n {
		this.c[x] += delta
		x += this.lowbit(x)
	}
}

func (this *BinaryIndexedTree) query(x int) int {
	s := 0
	x += 1e5 + 1
	for x > 0 {
		s += this.c[x]
		x -= this.lowbit(x)
	}
	return s
}

func subarraysWithMoreZerosThanOnes(nums []int) int {
	n := len(nums)
	s := make([]int, n+1)
	for i, v := range nums {
		if v == 0 {
			v = -1
		}
		s[i+1] = s[i] + v
	}
	tree := newBinaryIndexedTree(n + 1)
	ans := 0
	mod := int(1e9 + 7)
	for _, v := range s {
		ans = (ans + tree.query(v-1)) % mod
		tree.update(v, 1)
	}
	return ans
}

...