Add problem 3335: Total Characters in String After Transformations I

Author: EFanZh

src/lib.rs

@@ -2195,6 +2195,7 @@ pub mod problem_3326_minimum_division_operations_to_make_array_non_decreasing;
 pub mod problem_3330_find_the_original_typed_string_i;
 pub mod problem_3332_maximum_points_tourist_can_earn;
 pub mod problem_3334_find_the_maximum_factor_score_of_array;
+pub mod problem_3335_total_characters_in_string_after_transformations_i;
 pub mod problem_3350_adjacent_increasing_subarrays_detection_ii;
 
 #[cfg(test)]

src/problem_3335_total_characters_in_string_after_transformations_i/dynamic_programming.rs

@@ -0,0 +1,49 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl Solution {
+    fn add(mut lhs: u32, rhs: u32) -> u32 {
+        const MODULUS: u32 = 1_000_000_007;
+
+        lhs += rhs;
+
+        lhs.checked_sub(MODULUS).unwrap_or(lhs)
+    }
+
+    pub fn length_after_transformations(s: String, t: i32) -> i32 {
+        let t = t.cast_unsigned() as usize;
+        let mut counts = [0_u32; 26];
+
+        for c in s.bytes() {
+            counts[usize::from(c) - usize::from(b'a')] += 1;
+        }
+
+        let mut cursor = 0_usize;
+
+        for _ in 0..t {
+            let next_cursor = cursor.checked_sub(1).unwrap_or(25);
+
+            counts[cursor] = Self::add(counts[cursor], counts[next_cursor]);
+            cursor = next_cursor;
+        }
+
+        counts.iter().copied().fold(0, Self::add).cast_signed()
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn length_after_transformations(s: String, t: i32) -> i32 {
+        Self::length_after_transformations(s, t)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3335_total_characters_in_string_after_transformations_i/matrix_multiplication.rs

@@ -0,0 +1,95 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::array;
+
+impl Solution {
+    const MODULUS: u64 = 1_000_000_007;
+
+    fn matrix_multiply(lhs: &[[u32; 26]; 26], rhs: &[[u32; 26]; 26]) -> [[u32; 26]; 26] {
+        let mut result = [[0; 26]; 26];
+
+        result.iter_mut().zip(lhs).for_each(|(target_row, lhs_row)| {
+            for (i, target) in target_row.iter_mut().enumerate() {
+                *target = lhs_row.iter().zip(rhs).fold(0, |sum, (&lhs_value, rhs_row)| {
+                    (sum + u64::from(lhs_value) * u64::from(rhs_row[i])) % Self::MODULUS
+                }) as _;
+            }
+        });
+
+        result
+    }
+
+    fn exp_mod(mut base: [[u32; 26]; 26], mut exponent: u32) -> [[u32; 26]; 26] {
+        while exponent & 1 == 0 {
+            base = Self::matrix_multiply(&base, &base);
+            exponent >>= 1;
+        }
+
+        if exponent == 1 {
+            return base;
+        }
+
+        let mut result = base;
+
+        loop {
+            exponent >>= 1;
+            base = Self::matrix_multiply(&base, &base);
+
+            if exponent & 1 != 0 {
+                result = Self::matrix_multiply(&result, &base);
+
+                if exponent == 1 {
+                    break;
+                }
+            }
+        }
+
+        result
+    }
+
+    pub fn length_after_transformations(s: String, t: i32) -> i32 {
+        let base = array::from_fn(|i| {
+            let mut result = [0; 26];
+
+            result[25] = u32::from(i < 2);
+
+            if i != 0 {
+                result[i - 1] = 1;
+            }
+
+            result
+        });
+
+        let matrix = Self::exp_mod(base, t.cast_unsigned());
+
+        let mut counts = [0_u32; 26];
+
+        for c in s.bytes() {
+            counts[usize::from(c) - usize::from(b'a')] += 1;
+        }
+
+        matrix.iter().fold(0, |sum, row| {
+            row.iter().zip(&counts).fold(sum, |sum, (&lhs, &rhs)| {
+                (sum + u64::from(lhs) * u64::from(rhs)) % Self::MODULUS
+            })
+        }) as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn length_after_transformations(s: String, t: i32) -> i32 {
+        Self::length_after_transformations(s, t)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3335_total_characters_in_string_after_transformations_i/mod.rs

@@ -0,0 +1,19 @@
+pub mod dynamic_programming;
+pub mod matrix_multiplication;
+
+pub trait Solution {
+    fn length_after_transformations(s: String, t: i32) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [(("abcyy", 2), 7), (("azbk", 1), 5)];
+
+        for ((s, t), expected) in test_cases {
+            assert_eq!(S::length_after_transformations(s.to_string(), t), expected);
+        }
+    }
+}