Add problem 3186: Maximum Total Damage With Spell Casting

Author: EFanZh

src/lib.rs

@@ -2126,6 +2126,7 @@ pub mod problem_3178_find_the_child_who_has_the_ball_after_k_seconds;
 pub mod problem_3179_find_the_n_th_value_after_k_seconds;
 pub mod problem_3184_count_pairs_that_form_a_complete_day_i;
 pub mod problem_3185_count_pairs_that_form_a_complete_day_ii;
+pub mod problem_3186_maximum_total_damage_with_spell_casting;
 pub mod problem_3350_adjacent_increasing_subarrays_detection_ii;
 
 #[cfg(test)]

src/problem_3186_maximum_total_damage_with_spell_casting/dynamic_programming.rs

@@ -0,0 +1,72 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl Solution {
+    fn iter_groups(powers: &[u32], mut f: impl FnMut(u32, u64)) {
+        let mut prev_power = 0_u32;
+        let mut prev_power_sum = 0_u64;
+
+        for &power in powers {
+            if power == prev_power {
+                prev_power_sum += u64::from(power);
+            } else {
+                f(prev_power, prev_power_sum);
+
+                prev_power = power;
+                prev_power_sum = u64::from(power);
+            }
+        }
+
+        f(prev_power, prev_power_sum);
+    }
+
+    pub fn maximum_total_damage(power: Vec<i32>) -> i64 {
+        let mut power = power.into_iter().map(i32::cast_unsigned).collect::<Vec<_>>();
+
+        power.sort_unstable();
+
+        let mut power_queue = [0_u32; 4];
+        let mut damage_queue = [0_u64; 4];
+        let mut cursor = 0;
+
+        Self::iter_groups(&power, |power, damage| {
+            let diff = power.wrapping_sub(power_queue[cursor & 3]);
+
+            let offset = if diff < 3 {
+                if diff == 1 && power.wrapping_sub(power_queue[(cursor + 3) & 3]) == 2 {
+                    2
+                } else {
+                    3
+                }
+            } else {
+                0
+            };
+
+            let damage = damage_queue[cursor & 3].max(damage_queue[(cursor + offset) & 3] + damage);
+
+            cursor += 1;
+
+            power_queue[cursor & 3] = power;
+            damage_queue[cursor & 3] = damage;
+        });
+
+        damage_queue[cursor & 3].cast_signed()
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn maximum_total_damage(power: Vec<i32>) -> i64 {
+        Self::maximum_total_damage(power)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3186_maximum_total_damage_with_spell_casting/mod.rs

@@ -0,0 +1,22 @@
+pub mod dynamic_programming;
+
+pub trait Solution {
+    fn maximum_total_damage(power: Vec<i32>) -> i64;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            (&[1, 1, 3, 4] as &[_], 6),
+            (&[7, 1, 6, 6], 13),
+            (&[5, 9, 2, 10, 2, 7, 10, 9, 3, 8], 31),
+        ];
+
+        for (power, expected) in test_cases {
+            assert_eq!(S::maximum_total_damage(power.to_vec()), expected);
+        }
+    }
+}