3479. 水果成篮 III

src/main.rs

@@ -1,35 +1,69 @@
-use std::cmp::max;
+struct SegmentTree {
-use std::mem::swap;
+    tree: Vec<i32>,
+    arr: Vec<i32>,
+}
 
+impl SegmentTree {
+    pub fn new(baskets: Vec<i32>) -> Self {
+        let n = baskets.len();
+        let size = 2 << (32 - (n as u32 - 1).leading_zeros());
+        let mut seg = SegmentTree {
+            tree: vec![0; size as usize],
+            arr: baskets.clone(),
+        };
+        seg.build(1, 0, n - 1);
+        seg
+    }
+
+    fn push_up(&mut self, o: usize) {
+        self.tree[o] = self.tree[o * 2].max(self.tree[o * 2 + 1]);
+    }
+
+    fn build(&mut self, o: usize, l: usize, r: usize) {
+        if l == r {
+            self.tree[o] = self.arr[l];
+            return;
+        }
+        let m = (l + r) / 2;
+        self.build(o * 2, l, m);
+        self.build(o * 2 + 1, m + 1, r);
+        self.push_up(o);
+    }
+
+    pub fn find_update(&mut self, o: usize, l: usize, r: usize, x: i32) -> Option<usize> {
+        if self.tree[o] < x {
+            return None;
+        }
+        if l == r {
+            self.tree[o] = -1;
+            return Some(l);
+        }
+        let m = (l + r) / 2;
+        let mut i = self.find_update(o * 2, l, m, x);
+        if i.is_none() {
+            i = self.find_update(o * 2 + 1, m + 1, r, x);
+        }
+        self.push_up(o);
+        i
+    }
+}
 struct Solution;
 impl Solution {
-    pub fn total_fruit(fruits: Vec<i32>) -> i32 {
+    pub fn num_of_unplaced_fruits( fruits: Vec<i32>, baskets: Vec<i32>) -> i32 {
+        let n = baskets.len();
+        let mut tree = SegmentTree::new(baskets);
         let mut ans = 0;
-        let (mut a,mut b) = (-1,-1);
+
-        let mut temp_ans = 0;
+        for fruit in fruits.iter() {
-        let mut sub_temp_ans = 0;
+            if tree.find_update(1,0,n-1,*fruit).is_none() {
-        for fruit in fruits {
+                ans += 1;
-            if fruit != a && fruit != b {
-                ans = max(ans, temp_ans);
-                temp_ans = sub_temp_ans + 1;
-                sub_temp_ans = 1;
-                b = a;
-                a = fruit;
-            }
-            else {
-                temp_ans += 1;
-                sub_temp_ans += 1;
-                if b == fruit {
-                    swap(&mut a, &mut b);
-                    sub_temp_ans = 1;
-                }
             }
         }
-        max(ans, temp_ans)
+        ans
     }
 }
 
 fn main() {
-    let sl = Solution::total_fruit(vec![3,3,3,1,2,1,1,2,3,3,4]);
+    let sl = Solution::num_of_unplaced_fruits(vec![4,2,5],vec![3,5,4]);
     println!("{:?}", sl);
 }