Fixed string compression

tim-harding · Sep 1, 2023 · b19b135 · b19b135
1 parent c23b9e1
commit b19b135
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 18 deletions.
diff --git a/src/problems/_01_arrays_and_strings/_06_string_compression.rs b/src/problems/_01_arrays_and_strings/_06_string_compression.rs
@@ -26,10 +26,16 @@ mod tests {
     #[test]
     fn compresses() {
         assert_eq!(Solution::string_compression("aabcccccaaa"), "a2b1c5a3");
+        assert_eq!(Solution::string_compression("aaaaaaaaaaaa"), "a12");
+        assert_eq!(
+            Solution::string_compression(format!("{:a<1$}", "", 105).as_str()),
+            "a105"
+        );
     }
 
     #[test]
     fn returns_original_string() {
         assert_eq!(Solution::string_compression("abca"), "abca");
+        assert_eq!(Solution::string_compression("aabbcca"), "aabbcca");
     }
 }
diff --git a/src/solutions/_01_arrays_and_strings/_06_string_compression.rs b/src/solutions/_01_arrays_and_strings/_06_string_compression.rs
@@ -10,32 +10,40 @@ impl StringCompression for Solution {
             Some(c) => c,
             None => return Cow::Borrowed(string),
         };
-        loop {
-            match chars.next() {
-                Some(c) => {
-                    if c == previous {
-                        break;
-                    }
-                    previous = c;
-                }
-                None => return Cow::Borrowed(string),
-            }
-        }
-
-        let mut chars = string.chars();
-        let mut previous = chars.next().unwrap();
-        let mut count = 1u32;
+        let mut count = 1usize;
         let mut out = String::new();
         for c in chars {
             if c != previous {
-                out.extend(format!("{previous}{count}").chars());
+                out.push(previous);
+                append_count(&mut out, count);
                 previous = c;
                 count = 1;
             } else {
                 count += 1;
             }
         }
-        out.extend(format!("{previous}{count}").chars());
-        Cow::Owned(out)
+        out.push(previous);
+        append_count(&mut out, count);
+        if out.len() > string.len() {
+            Cow::Borrowed(string)
+        } else {
+            Cow::Owned(out)
+        }
+    }
+}
+
+fn append_count(s: &mut String, mut count: usize) {
+    // SAFETY: We are just appending ASCII characters, so this will yield valid
+    // UTF-8. We do this to avoid allocating with the format!() macro.
+    let v = unsafe { s.as_mut_vec() };
+    let mut count_chars = 0;
+    while count > 0 {
+        v.push((count % 10) as u8 + '0' as u8);
+        count /= 10;
+        count_chars += 1;
+    }
+    for i in 0..count_chars / 2 {
+        let len = v.len();
+        v.swap(len - count_chars + i, len - i - 1)
     }
 }