0006_Z字形变换

from itertools import chain

class Solution:
    def convert(self, s: str, numRows: int) -> str:
        if numRows == 1:
            return s
        g = [[] for _ in range(numRows)]
        i, k = 0, -1
        for c in s:
            g[i].append(c)
            if i == 0 or i == numRows - 1:
                k = -k
            i += k
        return ''.join(chain(*g))


if __name__ == "__main__":
    sol = Solution()
    test_cases = [
        ("PAYPALISHIRING", 3),
        ("PAYPALISHIRING", 4),
        ("A", 1),
        ("AB", 1),
        ("ABCDE", 2)
    ]
    for s, rows in test_cases:
        res = sol.convert(s, rows)
        print(f"s={repr(s)}, rows={rows} → {repr(res)}")