Add problem 2106: Maximum Fruits Harvested After at Most K Steps

src/lib.rs

@@ -1570,6 +1570,7 @@ pub mod problem_2102_sequentially_ordinal_rank_tracker;
 pub mod problem_2103_rings_and_rods;
 pub mod problem_2104_sum_of_subarray_ranges;
 pub mod problem_2105_watering_plants_ii;
+pub mod problem_2106_maximum_fruits_harvested_after_at_most_k_steps;
 pub mod problem_2108_find_first_palindromic_string_in_the_array;
 pub mod problem_2109_adding_spaces_to_a_string;
 pub mod problem_2110_number_of_smooth_descent_periods_of_a_stock;

src/problem_2106_maximum_fruits_harvested_after_at_most_k_steps/greedy.rs

@@ -0,0 +1,101 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl Solution {
+    pub fn max_total_fruits(fruits: Vec<Vec<i32>>, start_pos: i32, k: i32) -> i32 {
+        let fruits = fruits
+            .into_iter()
+            .map(|fruit| {
+                let [position, amount]: [_; 2] = fruit.try_into().ok().unwrap();
+
+                (position as u32, amount as u32)
+            })
+            .collect::<Vec<_>>();
+
+        let start_pos = start_pos as u32;
+        let k = k as u32;
+        let start_index = fruits.partition_point(|&(position, _)| position < start_pos);
+        let mut result = 0;
+
+        // Left.
+
+        let mut i = start_index;
+        let mut sum = 0;
+
+        loop {
+            i = i.wrapping_sub(1);
+
+            if let Some(&(position, amount)) = fruits.get(i) {
+                if position + k < start_pos {
+                    break;
+                }
+
+                sum += amount;
+            } else {
+                break;
+            }
+        }
+
+        let mut j = start_index;
+
+        for &(position, amount) in &fruits[i.wrapping_add(1)..start_index] {
+            while let Some(&(right_position, right_amount)) = fruits.get(j) {
+                if start_pos - position + (right_position - start_pos) * 2 <= k {
+                    sum += right_amount;
+                    j += 1;
+                } else {
+                    break;
+                }
+            }
+
+            result = u32::max(result, sum);
+            sum -= amount;
+        }
+
+        // Right.
+
+        while let Some(&(position, amount)) = fruits.get(j) {
+            if start_pos + k < position {
+                break;
+            }
+
+            sum += amount;
+            j += 1;
+        }
+
+        i = start_index.wrapping_sub(1);
+
+        for &(position, amount) in fruits[start_index..j].iter().rev() {
+            while let Some(&(left_position, left_amount)) = fruits.get(i) {
+                if position - start_pos + (start_pos - left_position) * 2 <= k {
+                    sum += left_amount;
+                    i = i.wrapping_sub(1);
+                } else {
+                    break;
+                }
+            }
+
+            result = u32::max(result, sum);
+            sum -= amount;
+        }
+
+        result as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn max_total_fruits(fruits: Vec<Vec<i32>>, start_pos: i32, k: i32) -> i32 {
+        Self::max_total_fruits(fruits, start_pos, k)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_2106_maximum_fruits_harvested_after_at_most_k_steps/mod.rs

@@ -0,0 +1,25 @@
+pub mod greedy;
+
+pub trait Solution {
+    fn max_total_fruits(fruits: Vec<Vec<i32>>, start_pos: i32, k: i32) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((&[[2, 8], [6, 3], [8, 6]] as &[_], 5, 4), 9),
+            ((&[[0, 9], [4, 1], [5, 7], [6, 2], [7, 4], [10, 9]], 5, 4), 14),
+            ((&[[0, 3], [6, 4], [8, 5]], 3, 2), 0),
+        ];
+
+        for ((fruits, start_pos, k), expected) in test_cases {
+            assert_eq!(
+                S::max_total_fruits(fruits.iter().map(Vec::from).collect(), start_pos, k),
+                expected,
+            );
+        }
+    }
+}