Count subarrays with more ones than zeros

Array Hash Table Binary Search Divide and Conquer Binary Indexed Tree Segment Tree Merge Sort Ordered Set

49.2% acceptance

Mar 31, 2026

187

No description available.

Solution

Rust

Time O(n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn subarrays_with_more_ones_than_zeroes(nums: Vec<i32>) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = nums.len();
    let size = 2 * n + 2;
    let mut bit = vec![0i64; size + 2];

    let update = |bit: &mut Vec<i64>, idx: usize| {
      let mut i = idx + 1;
      while i <= size {
        bit[i] += 1;
        i += i & i.wrapping_neg();
      }
    };

    let query = |bit: &Vec<i64>, idx: usize| -> i64 {
      let mut i = idx + 1;
      let mut s = 0i64;
      while i > 0 {
        s += bit[i];
        i -= i & i.wrapping_neg();
      }
      s
    };

    let offset = n;
    let mut prefix: i64 = 0;
    let mut ans: i64 = 0;

    update(&mut bit, offset);

    for &x in &nums {
      prefix += if x == 1 { 1 } else { -1 };
      let mapped = (prefix + offset as i64) as usize;
      if mapped > 0 {
        ans = (ans + query(&bit, mapped - 1)) % MOD;
      }
      update(&mut bit, mapped);
    }

    ans as i32
  }
}