fn main() {
    println!("Hello, world!");
}

pub fn convert(s: String, num_rows: i32) -> String {
    let bytes = s.as_bytes();
    let mut result = vec![vec![]; num_rows as usize];

    for (idx, byte) in bytes.iter().enumerate() {
        let row = if num_rows == 1 {
            0
        } else {
            let cycle_pos = idx % (2 * (num_rows - 1)) as usize;
            if cycle_pos < num_rows as usize - 1 {
                cycle_pos
            } else {
                2 * (num_rows as usize - 1) - cycle_pos
            }
        };

        result[row].push(*byte);
    }

    result.into_iter().flatten().map(|b| b as char).collect()
}

#[cfg(test)]
mod tests {
    use super::convert;

    #[test]
    fn test_1() {
        assert_eq!(
            convert("PAYPALISHIRING".to_string(), 3),
            "PAHNAPLSIIGYIR".to_string()
        );
    }
    #[test]
    fn test_2() {
        assert_eq!(
            convert("PAYPALISHIRING".to_string(), 4),
            "PINALSIGYAHRPI".to_string()
        );
    }
    #[test]
    fn test_3() {
        assert_eq!(convert("A".to_string(), 1), "A".to_string());
    }
    #[test]
    fn test_4() {
        assert_eq!(convert("AB".to_string(), 1), "AB".to_string());
    }
}