Improve Zig Zag conversion

src/leetcode/zigzag_conversion.rs

@@ -1,32 +1,42 @@
-fn add(u: usize, i: i32) -> Option<usize> {
-    if i.is_negative() {
-        u.checked_sub(i.wrapping_abs() as u32 as usize)
-    } else {
-        u.checked_add(i as usize)
-    }
-}
-
-pub fn convert(s: String, n: i32) -> String {
-    if n == 1 || n as usize >= s.len() {
-        return s;
+pub fn convert(s: String, num_rows: i32) -> String {
+    fn append(o: &mut String, bytes: &[u8], start: usize, step: usize) {
+        let mut i = 0;
+        while let Some(&c) = bytes.get(start + i * step) {
+            o.push(c as char);
+            i += 1;
+        }
     }
 
-    let mut rows = vec![String::new(); n as usize];
-
-    let mut row_index = 0;
-    let mut step = 1;
-    for c in s.chars() {
-        rows[row_index].push(c);
-        if row_index == 0 {
-            step = 1;
-        } else if row_index == n as usize - 1 {
-            step = -1;
+    let num_rows = num_rows as usize;
+    let bytes = s.as_bytes();
+    let step = match num_rows {
+        1 => 1,
+        n => n * 2 - 2,
+    };
+    let mut o = String::with_capacity(s.len());
+
+    // Top row
+    append(&mut o, bytes, 0, step);
+
+    // Middle rows
+    for row in 1..num_rows.saturating_sub(1) {
+        let mut i = 0;
+        while let Some(&a) = bytes.get(row + i * step) {
+            o.push(a as char);
+            match bytes.get((step - row) + i * step).copied() {
+                Some(b) => o.push(b as char),
+                None => break,
+            }
+            i += 1;
         }
+    }
 
-        row_index = add(row_index, step).unwrap();
+    // Bottom row
+    if num_rows > 1 {
+        append(&mut o, bytes, num_rows - 1, step);
     }
 
-    rows.join("")
+    o
 }
 
 #[cfg(test)]
@@ -43,6 +53,8 @@ mod tests {
         assert_eq!(
             "PINALSIGYAHRPI".to_string(),
             convert("PAYPALISHIRING".to_string(), 4)
-        )
+        );
+
+        assert_eq!("A".to_string(), convert("A".to_string(), 1));
     }
 }