diff --git a/src/main.rs b/src/main.rs index 4de2070..e62155f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,20 +1,40 @@ -use std::collections::{HashSet}; - -mod arr; -mod list; +use std::cmp::max; struct Solution; impl Solution { - pub fn find_final_value(nums: Vec, original: i32) -> i32 { - let mut original = original; - let set = nums.into_iter().collect::>(); - while set.contains(&original) { - original *= 2; + pub fn number_of_substrings(s: String) -> i32 { + let s = s.chars().collect::>(); + let mut ans = 0; + let mut zeros: Vec = Vec::new(); + + for i in 0..s.len() { + if s[i] == '0' { + zeros.push(i); + } + let mut local_ans =( if zeros.len() == 0 { + i + 1 + } else { + i - zeros[zeros.len() - 1] + }) as i32; + let mut ones = local_ans; + let zeros_len = zeros.len()as i32; + for l in (max(0, zeros_len - 200) ..zeros_len).rev() { + let l_index = l as usize; + let zeros_l_index = zeros[l_index] as i32; + let zero_start_index = if l == 0 { -1 } else { zeros[l_index - 1] as i32 }; + let need_ones = max(0, (zeros_len - l) * (zeros_len - l) - ones); + let max_index = zeros_l_index - need_ones; + ones += zeros_l_index - zero_start_index - 1; + if max_index < zero_start_index { + continue; + } + local_ans += max_index - zero_start_index; + } + ans += local_ans; } - original + ans } } - fn main() { - let result = Solution::find_final_value(vec![5,3,6,1,12],3); + let result = Solution::number_of_substrings("0001000011011010010011000100110100110101111111010011000101010110101000".to_string()); println!("{:?}", result); }