leetcode: 6. ZigZag Conversion

# The string "PAYPALISHIRING" is written in a zigzag pattern on a given number of rows like this: 
# (you may want to display this pattern in a fixed font for better legibility)
# 
# P   A   H   N
# A P L S I I G
# Y   I   R
# And then read line by line: "PAHNAPLSIIGYIR"
# Write the code that will take a string and make this conversion given a number of rows:
# 
# string convert(string text, int nRows);
# convert("PAYPALISHIRING", 3) should return "PAHNAPLSIIGYIR".

# Time:  O(n)
# Space: O(1)
class Solution():
    def convert(self, s, numRows):
        if numRows == 1:
            return s
        step, zigzag = 2 * numRows - 2, ""
        for i in range(numRows):
            for j in range(i, len(s), step):
                zigzag += s[j]
                if 0 < i < numRows - 1 and j + step - 2 * i < len(s):
                    zigzag += s[j + step - 2 * i]
        return zigzag


if __name__ == "__main__":
    assert Solution().convert("PAYPALISHIRING", 1) == 'PAYPALISHIRING'
    assert Solution().convert("PAYPALISHIRING", 3) == 'PAHNAPLSIIGYIR'
    assert Solution().convert("PAYPALISHIRING", 4) == 'PINALSIGYAHRPI'