2106. 摘水果

src/main.rs

@@ -1,54 +1,44 @@
-use std::cmp::min;
-use std::collections::HashMap;
+use std::cmp::max;
 
 struct Solution;
 impl Solution {
-    pub fn min_cost(basket1: Vec<i32>, basket2: Vec<i32>) -> i64 {
-        let mut map : HashMap<i32,i32> = HashMap::new();
-        let mut total_cost: i64 = 0;
-        let mut min_cost = i32::MAX;
-        for i in basket1.iter() {
-            *map.entry(*i).or_insert(0) += 1;
+    pub fn max_total_fruits(fruits: Vec<Vec<i32>>, start_pos: i32, k: i32) -> i32 {
+        let mut left = fruits
+            .binary_search_by(|x| x[0].cmp(&(start_pos - k)))
+            .unwrap_or_else(|x| x);
+        let (mut right, mut temp_ans) = (left, 0);
+        while right < fruits.len() && fruits[right][0] <= start_pos {
+            temp_ans += fruits[right][1];
+            right += 1;
         }
-        for i in basket2.iter() {
-            *map.entry(*i).or_insert(0) -= 1;
-        }
-        let mut kvs: Vec<(i32, i32)> = map.iter().filter(|(k, v)| {min_cost = min(min_cost, **k); if **v == 0 {false} else {true}}).map(|(k, v)| (*k, *v)).collect();
-        let mut err = false;
-        kvs.sort_by(|a, b| {
-            err |= b.1 % 2 == 1 || a.1 % 2 == 1;
-            if a.1 > 0 && b.1 < 0 {
-                std::cmp::Ordering::Greater
-            } else if a.1 < 0 && b.1 > 0 {
-                std::cmp::Ordering::Less
-            } else {
-                a.0.cmp(&b.0)
+        let mut ans = temp_ans;
+        while right < fruits.len() && fruits[right][0] <= start_pos + k {
+            temp_ans += fruits[right][1];
+            while fruits[right][0] - fruits[left][0] + fruits[right][0] - start_pos > k
+                && fruits[right][0] - fruits[left][0] + start_pos - fruits[left][0] > k
+            {
+                temp_ans -= fruits[left][1];
+                left += 1;
             }
-        });
-        if err
-        { return -1; }
-        if kvs.is_empty() {
-            return 0;
+            ans = max(ans, temp_ans);
+            right += 1;
         }
-        let mut l = 0;
-        let mut r = kvs.len() - 1;
-        while l < r {
-            let cost = min( kvs[l].1.abs(), kvs[r].1.abs()) as i64;
-            total_cost += min(min(kvs[l].0, kvs[r].0) as i64 * cost / 2,min_cost as i64 * cost);
-            kvs[l].1 += cost as i32;
-            kvs[r].1 -= cost as i32;
-            if kvs[l].1 >= 0 {
-                l+=1;
-            }
-            if kvs[r].1 <= 0 {
-                r-=1;
-            }
-        }
-        if kvs[l].1 > 0 || kvs[r].1 > 0 { -1 } else { total_cost }
+        ans
     }
 }
 
 fn main() {
-    let sl = Solution::min_cost(vec![4,2,2,2], vec![1,4,1,2]);
+    let sl = Solution::max_total_fruits(
+        vec![
+            vec![0, 9],
+            vec![4, 1],
+            vec![5, 7],
+            vec![6, 2],
+            vec![7, 4],
+            vec![10, 9],
+        ],
+        5,
+        4,
+    );
     println!("{:?}", sl);
 }