LeetCode 6. Zigzag Conversion

A super naive $O(n^2)$ solution would be to make an $n \times n$ matrix, actually produce the zigzag image, and then parse it out.

But I think there's probably a mathematical solution. I'm also going to guess that it's linear & based on index offsets.

The first character is always $s[0]$.

The columns of the zigzag are of length $numRows$ and the diagonals are of length $numRows - 2$.

So the second character is at index $0 + numRows + numRows-2$. We can use this offset pattern to get every column in every row.

Now for the diagonals.

In row $0$, there are no diagonals. In row $1$, the diagonals are at the location of the current column plus 4. In row $2$, the diagonals are at the location of the current column plus 2, and so on. In row $numRows-1$, there are no diagonals a.k.a. the diagonals would compute to be at the same index as the current column (this is similar to why there is no diagonal for the first row).

So we can iterate over rows, alternating between column and diagonal calculations, and stopping each row when we exceed the bounds of $s$.

The runtime of this algorithm is $O(n)$ and the space complexity is $O(1)$.

import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_leetcode_examples(self):
        self.assertEqual(self.s.convert("PAYPALISHIRING", 3), "PAHNAPLSIIGYIR")
        self.assertEqual(self.s.convert("PAYPALISHIRING", 4), "PINALSIGYAHRPI")
        self.assertEqual(self.s.convert("A", 1), "A")
        self.assertEqual(self.s.convert("AB", 1), "AB")

class Solution:

    def convert(self, s: str, numRows: int) -> str:
        if numRows == 1:
            return s
        res = ""
        for i in range(numRows):
            col = i
            while col < len(s):
                res += s[col]
                col_offset = 2 * (numRows - 1)
                col_next =  col + col_offset
                diag = col + 2 * (numRows - 1 - col % col_offset)
                if diag != col_next and diag != col and diag < len(s):
                    res += s[diag]
                col = col_next
        return res
                    


if __name__=="__main__":
    unittest.main()