3501. 操作后最大活跃区段数 II

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;
+    fn get_ans_from_tree(tree: &Vec<i32>, 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<i32>,
+        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<i32>>) -> Vec<i32> {
+        let arr = s
+            .chars()
+            .map(|x| -> u8 { x as u8 - b'0' })
+            .collect::<Vec<u8>>();
+        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 {
-                ans += 1;
-                if zero_count > 0{
-                    if last_zero_count > 0 {
-                        max_p = max_p.max(last_zero_count + zero_count);
+                    zero_count = zero_count + 1;
                 }
-                    last_zero_count = zero_count;
-                    zero_count = 0
+            } 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]]")
+        )
+    );
 }