Zig-zag function - coded solution

Question

Here is my coded solution to the LeetCode zig-zag problem -

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.

I know the final line is a bit messy, and I could unpack it into a two-line for loop, but I really like single-line for loops.

def zigzag() -> list:
    """Returns the zig-zag word of a given input as a list."""
    array = []
    for _ in range(num_rows):
        array.append([])
    increasing = 1
    array_index = 0
    word_index = 0
    while word_index < len(word):
        if increasing:
            array[array_index].append(word[word_index])
            array_index += 1
            word_index += 1
        else:
            array[array_index].append(word[word_index])
            array_index -= 1
            word_index += 1
        if array_index == -1 and increasing == 0:
            increasing = 1
            array_index = 1
        if array_index == num_rows and increasing == 1:
            increasing = 0
            array_index = num_rows - 2
    return array

if __name__ == "__main__":
    word = "PAYPALISHIRING"
    num_rows = 3
    print("".join(line for array in zigzag() for line in array))

Answer 1

PEP 8 implementation is great, now to focus on being concise -

While you have started writing immaculate programs, as in here and your current question, I believe you should start making your programs more concise. This decreases memory space and sometimes, makes your program faster. Also, I'm glad to see you've started implementing the if __name__ == '__main__' guard.

Here's one approach you could use to make your program really concise and fast -

def zigzag(word: str, num_rows: int) -> str:

    ans = [''] * num_rows
    x = 0
    direction = 1
    for i in word:
        ans[x] += i
        if 0 <= x + direction < num_rows:
             x += direction
        else:
             direction *= -1
             x += direction
    return ''.join(ans)

if __name__ == "__main__":
    print(zigzag("PAYPALISHIRING", 3))

Here, I believe that it would be better if you don't declare word and num_rows before calling the function. This -

print(zigzag("PAYPALISHIRING", 3))

looks much shorter (and better) than -

word = "PAYPALISHIRING"
num_rows = 3
print("".join(line for array in zigzag() for line in array))

Also, (hey, I'm kind of like a beginner, so please correct me if I'm wrong) I feel that you should perform the .join() function, or any other function really, in the principal function, as I have done above -

# rest of the principal function
return ''.join(ans)

Just use the if __name__ == __main__ guard to (simply) call the principal function and execute the entire program (you shouldn't really perform any other function(s) in it except calling the principal function to execute it) -

if __name__ == "__main__":
    print(zigzag("PAYPALISHIRING", 3))

---

Also, I would rename increasing to direction as it suits this task better ("the value of "direction" changes only when we have moved up to the topmost row or moved down to the bottommost row" (taken from Leetcode). If you're wondering what *= means, it's just another way of doing x = x * 5, which is specified here.

Also, you could make -

array = []
for _ in range(num_rows):
    array.append([])

num_rows = 3
print(array)
# [[], [], []]

more concise by just simply initializing it -

array = [[] for _ in range(num_rows)]  # Thanks to Maarten Fabré

num_rows = 3
print(array)
# [[], [], []]

See how concise you can get by doing this? This removes the need to use .append(), which, in this case, is unnecessary (this will also take up less memory space).

---

I know the final line is a bit messy, and I could unpack it into a two-line for loop but I really like single-line for loops.

Nah, don't worry. As far as I know,

print("".join(line for array in zigzag() for line in array))

is so much more concise than unpacking it into a two-line for loop. I'm glad you like single-line for loops :)

Here are the times taken for each function -

Your function -

%timeit "".join(line for array in zigzag() for line in array)
8.34 µs ± 834 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

---

My function -

%timeit zigzag("PAYPALISHIRING", 3)
4.19 µs ± 134 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

As you can see above, the time taken for execution has been approximately halved, which also shows that making your function concise, down to what I have done in my function, can make your program much faster.

Here is the Leetcode result for my program (if need be) -

Hope this helps!