diff --git a/src/main.rs b/src/main.rs index 137ad18..d77a465 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,30 +1,134 @@ +use crate::arr::make_matrix; + mod arr; struct Solution {} impl Solution { - pub fn max_active_sections_after_trade(s: String) -> i32 { - let (mut last_zero_count, mut zero_count) = (0,0); - let mut ans = 0; - let mut max_p = 0; - for c in s.chars() { - if c == '0' { - zero_count += 1; - }else { - ans += 1; - if zero_count > 0{ - if last_zero_count > 0 { - max_p = max_p.max(last_zero_count + zero_count); - } - last_zero_count = zero_count; - zero_count = 0 + fn get_ans_from_tree(tree: &Vec, layer_count: usize, l: usize, r: usize) -> usize { + let (cl, cr) = (0, 2_usize.pow(layer_count as u32 - 1) - 1); + Self::dfs(tree, 1, cl, cr, l, r) + } + fn dfs( + tree: &Vec, + root: usize, + l: usize, + r: usize, + target_l: usize, + target_r: usize, + ) -> usize { + if l >= target_l && r <= target_r { + return tree[root] as usize; + } + if l > target_r || r < target_l { + return 0; + } + let from_l = Self::dfs(tree, root * 2, l, (r - l) / 2 + l, target_l, target_r); + let from_r = Self::dfs( + tree, + root * 2 + 1, + (r - l) / 2 + l + 1, + r, + target_l, + target_r, + ); + from_l.max(from_r) + } + pub fn max_active_sections_after_trade(s: String, queries: Vec>) -> Vec { + let arr = s + .chars() + .map(|x| -> u8 { x as u8 - b'0' }) + .collect::>(); + let mut zero_start = vec![]; + let mut zero_end = vec![]; + let mut zero_len = vec![]; + // let mut one_count = vec![0]; + let mut one_count = 0; + let mut zero_count = -1; + for i in 0..arr.len() { + if arr[i] == 0 { + // one_count.push(one_count[one_count.len() - 1]); + if zero_count == -1 { + zero_start.push(i); + zero_count = 1; + } else { + zero_count = zero_count + 1; + } + } else { + // one_count.push(one_count[one_count.len() - 1] + 1); + one_count += 1; + if zero_count != -1 { + zero_len.push(zero_count); + zero_end.push(i - 1); + zero_count = -1; } } } - if zero_count > 0 && last_zero_count > 0 { - max_p = max_p.max(zero_count + last_zero_count); + if zero_count != -1 { + zero_len.push(zero_count); + zero_end.push(arr.len() - 1); } - ans + max_p + + let len = zero_len.len() as f64 - 1f64; + let layer_count = len.log2().ceil() as usize + 1; + let tree_len = 2_usize.pow(layer_count as u32); + let mut tree = vec![0; tree_len]; + let base = 2_usize.pow(layer_count as u32 - 1); + if zero_len.len() > 0 { + for i in 0..zero_len.len() - 1 { + tree[base + i] = zero_len[i] + zero_len[i + 1]; + } + for i in (1..base).rev() { + tree[i] = tree[i * 2].max(tree[i * 2 + 1]); + } + } + let mut ans = vec![]; + for query in queries { + let (l, r) = (query[0] as usize, query[1] as usize); + // let one = one_count[one_count.len() - 1] - one_count[0]; + let one = one_count; + let i = zero_start.partition_point(|&x| x < l); + let j = zero_end.partition_point(|&x| x <= r); + let mut local_ans = 0; + + if j as i32 - i as i32 >= 2 { + local_ans = local_ans.max(Self::get_ans_from_tree(&tree, layer_count, i, j - 2)); + } + if j as i32 - i as i32 >= 1 { + if i > 0 && i < zero_end.len() { + if zero_end[i - 1] >= l { + let pz = zero_end[i] - zero_start[i] + 1 + zero_end[i - 1] - l + 1; + local_ans = local_ans.max(pz); + } + } + if j < zero_end.len() && j > 0 { + if r >= zero_start[j] { + let pz = zero_end[j - 1] - zero_start[j - 1] + 1 + r - zero_start[j] + 1; + local_ans = local_ans.max(pz); + } + } + } + if i == j { + if i > 0 + && i < zero_len.len() + && j < zero_len.len() + && j > 0 + && zero_end[i - 1] >= l + && zero_start[i] <= r + { + let pz = zero_end[j - 1] - l + 1 + r - zero_start[j] + 1; + local_ans = local_ans.max(pz); + } + } + ans.push((local_ans + one) as i32); + } + ans } } fn main() { - println!("{:?}", Solution::max_active_sections_after_trade("0100".to_string())); + println!( + "{:?}", + Solution::max_active_sections_after_trade( + "0110110101010000101".to_string(), + make_matrix("[[12,17]]") + ) + ); }