main.py

#!/usr/bin/env python3

'''
Input: "MICROSOFTISHIRING", num_rows=3:
--------------------------------------
M   O   T   I   G
I R S F I H R N
C   O   S   I
=> Output: "MOTIGIRSFIHRNCOSI"

Input: "MICROSOFTISHIRING", num_rows=4:
--------------------------------------
M     O      I
I   S F    H R
C O   T  S   I  G
R     I      N
=> Output: "MOIISFHRCOTSIGRIN"

Input: "A", num_rows = 1:
------------------------
A
=> Output: "A"

Contraints:
(1) num_rows >= 1
(2) len(Input) >= 1
(3) Time complexity: O(n)
(4) Memory complexity: O(n)
'''

def zigzag(s: str, num_rows: int):
    # O(n) memory
    sol = ''
    # Amount of characters that will be saved is exactly len(s) -> O(n) memory
    rows = [[] for _ in range(num_rows)] # num_rows = 2: rows = [[], []]
    current_row = 0
    if num_rows == 1:
        next_row_dir = 0
    else:
        next_row_dir = -1

    # 0(n) time complexity
    for c in s:
        rows[current_row].append(c)
        if current_row == 0 or current_row == num_rows - 1:
            next_row_dir *= -1
        current_row += next_row_dir

    # O(n) time complexity
    for row in rows:
        for c in row:
            sol += c

    '''
    Time: O(n) + O(n) = O(n)
    Memory: O(n) + O(n) = O(n)
    '''
    return sol


def main():
    print(zigzag('MICROSOFTISHIRING', 3) == 'MOTIGIRSFIHRNCOSI')
    print(zigzag('ABCD', 1) == 'ABCD')

if __name__ == '__main__':
    main()