diff --git a/src/leetcode/zigzag_conversion.rs b/src/leetcode/zigzag_conversion.rs index 847e2ff..ac4bf2f 100644 --- a/src/leetcode/zigzag_conversion.rs +++ b/src/leetcode/zigzag_conversion.rs @@ -1,32 +1,42 @@ -fn add(u: usize, i: i32) -> Option { - 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)); } }