087_zigzag_conversion.sio

//@ run-pass
// HumanEval 087: Zigzag Conversion
//
// Rearrange characters in a zigzag pattern with a given number of rows,
// then read off row by row. For "PAYPALISHIRING" with 3 rows:
//   P   A   H   N
//   A P L S I I G
//   Y   I   R
// Result: "PAHNAPLSIIGYIR"
// Input/output as byte arrays.

struct Buf { data: [i64; 256] }

fn zigzag(input: [i64; 256], n: i64, num_rows: i64, out: &!Buf) -> i64 with Mut, Panic, Div {
    if num_rows <= 1 || num_rows >= n {
        // Just copy
        var k: i64 = 0
        while k < n {
            out.data[k] = input[k]
            k = k + 1
        }
        return n
    }

    // Cycle length = 2 * (num_rows - 1)
    let cycle = 2 * (num_rows - 1)
    var out_idx: i64 = 0

    var row: i64 = 0
    while row < num_rows {
        var i: i64 = row
        while i < n {
            out.data[out_idx] = input[i]
            out_idx = out_idx + 1

            // For non-first and non-last rows, there is a mid-cycle char
            if row > 0 && row < num_rows - 1 {
                let mid = i + cycle - 2 * row
                if mid < n {
                    out.data[out_idx] = input[mid]
                    out_idx = out_idx + 1
                }
            }

            i = i + cycle
        }
        row = row + 1
    }

    out_idx
}

fn main() -> i64 with IO, Mut, Panic, Div {
    // Test 1: "PAYPALISHIRING" (3 rows) -> "PAHNAPLSIIGYIR"
    // P=80 A=65 Y=89 P=80 A=65 L=76 I=73 S=83 H=72 I=73 R=82 I=73 N=78 G=71
    var s1: [i64; 256] = [0; 256]
    s1[0] = 80
    s1[1] = 65
    s1[2] = 89
    s1[3] = 80
    s1[4] = 65
    s1[5] = 76
    s1[6] = 73
    s1[7] = 83
    s1[8] = 72
    s1[9] = 73
    s1[10] = 82
    s1[11] = 73
    s1[12] = 78
    s1[13] = 71
    var o1 = Buf { data: [0; 256] }
    let n1 = zigzag(s1, 14, 3, &!o1)
    assert(n1 == 14)
    // Expected: P=80 A=65 H=72 N=78 A=65 P=80 L=76 S=83 I=73 I=73 G=71 Y=89 I=73 R=82
    assert(o1.data[0] == 80)
    assert(o1.data[1] == 65)
    assert(o1.data[2] == 72)
    assert(o1.data[3] == 78)
    assert(o1.data[4] == 65)
    assert(o1.data[5] == 80)
    assert(o1.data[6] == 76)
    assert(o1.data[7] == 83)
    assert(o1.data[8] == 73)
    assert(o1.data[9] == 73)
    assert(o1.data[10] == 71)
    assert(o1.data[11] == 89)
    assert(o1.data[12] == 73)
    assert(o1.data[13] == 82)

    // Test 2: "AB" with 1 row -> "AB" (no change)
    var s2: [i64; 256] = [0; 256]
    s2[0] = 65
    s2[1] = 66
    var o2 = Buf { data: [0; 256] }
    let n2 = zigzag(s2, 2, 1, &!o2)
    assert(n2 == 2)
    assert(o2.data[0] == 65)
    assert(o2.data[1] == 66)

    // Test 3: "ABCD" with 2 rows -> "ACBD"
    // Row 0: A C
    // Row 1: B D
    var s3: [i64; 256] = [0; 256]
    s3[0] = 65
    s3[1] = 66
    s3[2] = 67
    s3[3] = 68
    var o3 = Buf { data: [0; 256] }
    let n3 = zigzag(s3, 4, 2, &!o3)
    assert(n3 == 4)
    assert(o3.data[0] == 65)
    assert(o3.data[1] == 67)
    assert(o3.data[2] == 66)
    assert(o3.data[3] == 68)

    // Test 4: single char "A" with 3 rows -> "A"
    var s4: [i64; 256] = [0; 256]
    s4[0] = 65
    var o4 = Buf { data: [0; 256] }
    let n4 = zigzag(s4, 1, 3, &!o4)
    assert(n4 == 1)
    assert(o4.data[0] == 65)

    // Test 5: "ABCDE" with 4 rows
    // Row 0: A     -> A
    // Row 1: B   E -> B E
    // Row 2: C     -> C
    // Row 3: D     -> D
    // Result: "ABECD"
    var s5: [i64; 256] = [0; 256]
    s5[0] = 65
    s5[1] = 66
    s5[2] = 67
    s5[3] = 68
    s5[4] = 69
    var o5 = Buf { data: [0; 256] }
    let n5 = zigzag(s5, 5, 4, &!o5)
    assert(n5 == 5)
    assert(o5.data[0] == 65)
    assert(o5.data[1] == 66)
    assert(o5.data[2] == 69)
    assert(o5.data[3] == 67)
    assert(o5.data[4] == 68)

    println("087_zigzag_conversion: ALL TESTS PASSED")
    0
}