Zigzag and wave in Python

Why this matters

A zigzag is a path that goes one way, turns, and comes back. A ball bouncing between a floor and a ceiling, a wave drawn with characters, and text written in a zigzag all follow the same rule: the position moves one step at a time, and when it reaches an edge it turns around.

This lesson teaches that rule in Python. You will keep two variables, one for the current row and one for the direction, and a check that turns the direction around at the edges. The other lessons in this module print shapes whose characters come from a fixed formula. The shape in this lesson comes from a position that changes on every step, so the loop has to update a variable each time, not only print a character.

Some terms first. A grid is a table made of rows and columns. Rows run across, columns run down, and in this lesson row 0 is the top row. An index is a position number that starts at 0. A list is an ordered collection of values, written in square brackets, such as [3, 1, 4]. You reach an item by its index, and you can change an item in place. A grid in Python is a list of lists: the first index picks the row and the second picks the column, so grid[row][col] is the cell in that row and that column. The direction is a number that says where the marker moves next. Here 1 means one row down and -1 means one row up.

If you know Java, a Python list plays the part of a Java array, but it can grow and shrink, and its items can be of any type.

The idea: one marker, one direction

Start with a board of 4 rows. The rows are numbered 0 to 3 from the top. A marker starts on row 0 and moves one row per column. In each column we place a star on the marker's row. The stars form a zigzag.

The rule has two parts, and the order matters:

1. Look at the current row. If it is the top row (0), the next move must go down, so the direction becomes 1. If it is the bottom row (3), the next move must go up, so the direction becomes -1. Otherwise the direction stays as it was.
2. Move the marker: add the direction to the row.

Here is the trace for 12 columns. Each line is one column. "Row now" is where the star goes. "Direction" is the value after the check in part 1.

column  row now  direction  row after the move
0       0        1          1
1       1        1          2
2       2        1          3
3       3        -1         2
4       2        -1         1
5       1        -1         0
6       0        1          1
7       1        1          2
8       2        1          3
9       3        -1         2
10      2        -1         1
11      1        -1         0

The rows where stars land are 0 1 2 3 2 1 0 1 2 3 2 1. The marker goes down for three moves and up for three moves, so the pattern repeats every 6 columns.

If the check came after the move, the marker would step past the bottom edge before it turned. The check has to look at the row the marker is standing on now.

Example 1: a bouncing star on a grid

The first program draws the zigzag as stars on a 4 by 12 grid. It fills the grid with spaces, then walks across the columns, placing a star at each step.

The output is:

*     *     
 *   * *   *
  * *   * * 
   *     *  

Each line of the output has 12 characters, including the spaces at the end. Here is what each part does.

[" "] * width makes a list of 12 spaces. The next line, [[" "] * width for _ in range(rows)], makes 4 of those lists, one per row. The for inside the square brackets is a list comprehension, which is a short way to build a list with a loop. The name _ is a placeholder for a value we do not use. Each row is a separate list, which matters in the Common mistakes section below.

grid[row][col] = "*" places a star in the marker's current row, in the current column. Python lists can be changed in place, so this line changes the cell itself.

The if and elif are the check from part 1. They run before the move. elif means "else if". At the top row the direction is set to 1. At the bottom row, rows - 1 is 3, so the direction is set to -1. In the middle rows neither branch runs and the direction keeps its old value.

row += direction is the move. With direction 1 it adds one, which goes down. With -1 it subtracts one, which goes up.

"".join(line) joins the characters of one row into a single piece of text. The empty string before .join is the separator put between the pieces, so here nothing is put between them. print then writes that text on its own line.

The code sets the direction to a fixed value at each edge instead of reversing it with direction = -direction. Setting the value says what you mean: at the top, go down. A reversal depends on the old value being right. If the check ever ran twice on the same edge, a reversal would undo itself. The fixed value gives the same answer every time.

Example 2: a bounce that looks ahead

The first version turns around after it reaches an edge. Another way is to compute the next row first, check whether that row is still on the board, and turn around only if it is not. "On the board" means the row index is between 0 and rows - 1. A row index outside that range is out of bounds.

This version also records each row in a list, so you can see the full path.

The output is:

rows visited: [0, 1, 2, 3, 4, 3, 2, 1, 0, 1, 2, 3]

Trace the steps around the bottom edge. The marker is on row 4, which is the last row. The next row would be 5. Since 5 is not less than rows (5), the check next_row >= rows is true. The direction flips from 1 to -1, and the next row becomes 4 + -1 = 3. The marker is on row 4 at step 4 and on row 3 at step 5, which matches the output.

The variable is called next_row and not next, because next is a built-in function in Python, and using it as a name would hide the function.

This version reads as one sentence: "turn around when the next step would leave the board." It does not depend on remembering which edge is the top and which is the bottom. The same code works for any board with two or more rows, and any starting row inside the board, and it is a common way to keep a position inside a fixed range.

visited.append(row) adds the row to the end of the list. It runs before the move, so the first entry is the starting row, 0.

Example 3: text in a zigzag

The next problem uses the same movement to rearrange text. Take a string, write its letters in a zigzag over a fixed number of rows, then read the rows from top to bottom, one after another. The string PAYPALISHIRING with 3 rows is a standard example.

First, here is the zigzag as a picture. A function is a named block of code that you can call with values, and def starts one. Each letter goes in its own column, in the row where the marker is at that step.

The output is:

P   A   H   N 
 A P L S I I G
  Y   I   R   
 
PAHNAPLSIIGYIR
PINALSIGYAHRPI
AB

The picture shows the zigzag. Read down the first column: P is on row 0, A is on row 1, Y is on row 2. Read across row 1 and you get A, P, L, S, I, I, G. Those are the letters that the zigzag placed on the middle row, in the order they appear in the text.

The function row_of gives the row for any letter position without a loop. It uses a cycle, which is the number of steps before the pattern repeats. % is the remainder operator: 7 % 4 is 3. With 3 rows the cycle is 2 * (3 - 1), which is 4. Positions 0 to 2 go down from row 0 to row 2. Position 3 steps back up to row 1. In code, position is the index inside one cycle. If it is less than num_rows, the row is the position. Otherwise the row is cycle - position. For 3 rows the list of rows for positions 0 to 7 is 0 1 2 1 0 1 2 1. The function needs at least 2 rows. With 1 row the cycle is 0, and % by 0 is an error.

enumerate(text) gives each letter together with its position, so i is the position and ch is the letter. Python strings are read one character at a time in a for loop, as convert does.

The convert function does the same walk as Example 1, but it appends each letter to a list of lists, one list per row. After the walk, the last line joins the rows top to bottom. The inner "".join(line) joins the letters of one row. The outer "".join(...) joins the rows, one after another. The part in the brackets is a generator expression, which is a short way to write a loop that produces one piece at a time. The result is the text read row by row, which is what the problem asks for.

The last four lines check the two versions with 3 rows and 4 rows. For PAYPALISHIRING with 3 rows, the joined rows are PAHNAPLSIIGYIR. With 4 rows they are PINALSIGYAHRPI. A two-letter string with 2 rows gives AB, because each row has one letter.

Common mistakes

Turning at the edges with the wrong check

This version turns at the edges too early, because the check is wrong:

rows = 4
row = 0
direction = 1
for step in range(5):
    print("step", step, "row", row)
    if row == 0 or row == rows - 1:
        direction = -direction
    row += direction

The output is:

step 0 row 0
step 1 row -1
step 2 row -2
step 3 row -3
step 4 row -4

The check says "reverse at the top or the bottom". On the very first step the marker is on row 0, so the direction is reversed from 1 to -1 before the first move. The marker then walks up and off the board, into rows -1, -2, and so on.

Python does not stop here, and this is the surprise. Python lists accept negative indexes. grid[-1] means the last item, grid[-2] the one before it, and so on. The next program shows what that does with a 4-item list:

grid = ["row 0", "row 1", "row 2", "row 3"]
print(grid[-1])
print(grid[-4])
print(grid[-5])

The output is:

row 3
row 0
Traceback (most recent call last):
  File "bad_flip2.py", line 4, in <module>
    print(grid[-5])
          ~~~~^^^^
IndexError: list index out of range

So the wrong row numbers -1 to -4 quietly pick real rows from the bottom of the list, and no error appears. Only -5 fails. The file path and the markers under the code depend on where you saved the file and on your Python version. This is why a wrong turn can give a picture that looks almost right. The fix is to set the direction to a fixed value at each edge, as in Example 1. The top edge sets 1 and the bottom edge sets -1. The first move then goes down, as it should.

Forgetting the one-row board

The convert function breaks when num_rows is 1. Here is the same function, called with a board of one row:

def convert(text, num_rows):
    rows = [[] for _ in range(num_rows)]
    row = 0
    direction = 1
    for ch in text:
        rows[row].append(ch)
        if row == 0:
            direction = 1
        elif row == num_rows - 1:
            direction = -1
        row += direction
    return "".join("".join(line) for line in rows)
 
 
print(convert("AB", 1))

The output is:

Traceback (most recent call last):
  File "z_one.py", line 15, in <module>
    print(convert("AB", 1))
          ~~~~~~~^^^^^^^^^
  File "z_one.py", line 6, in convert
    rows[row].append(ch)
    ~~~~^^^^^
IndexError: list index out of range

With one row, the marker is always on row 0. The if row == 0 branch runs, which sets the direction to 1, and then row += direction moves the marker to row 1. The list has only one item, so index 1 fails. The last line says what went wrong: the index is out of range for the list. The line with rows[row].append(ch) shows where. The path and the line numbers depend on where you saved the file, so yours will differ. The fix is a special case at the top of convert:

if num_rows == 1:
    return text

With one row, the zigzag is a straight line, so the text is returned as it is.

Sharing one row between all rows

Building the grid with [[" "] * width] * rows looks shorter, and it is wrong:

width = 4
rows = 3
 
bad = [[" "] * width] * rows
bad[0][0] = "*"
for line in bad:
    print("".join(line))
 
print()
 
good = [[" "] * width for _ in range(rows)]
good[0][0] = "*"
for line in good:
    print("".join(line))

The output is:

*   
*   
*   
 
*   
    
    

The * operator makes rows references to the same list, so changing one row changes all of them. The fix is the list comprehension from Example 1, which builds a new list for each row.

Forgetting to move

If the line row += direction is missing, the marker never leaves row 0. The output is a straight line of stars on the top row. The direction changes do nothing because the row never changes. When a zigzag looks flat, check the move first.

Swapping the row and the column

In grid[row][col] the first index is the row. Writing grid[col][row] treats the columns as rows. The marker row goes into the wrong place. Once the column number passes the number of rows, the first index is out of range, and the program stops with an IndexError. The fix is to keep the order: row first, column second. A quick check is to ask "how many rows does this grid have?" and make sure that the first index never goes above that number minus 1.

Practice

Exercise 1: trace by hand

With 3 rows, start at row 0 with direction 1. Use the rule from Example 1. Write the row for each of the first 7 columns.

Show the answer

The rows are 0 1 2 1 0 1 2. The marker goes down to row 2, turns, goes up to row 0, and turns again. Column 3 is row 1 because the marker moved up from row 2 at column 2.

Exercise 2: change the board

Change rows to 3 and width to 9 in the Example 1 program. Write the output on paper first, then run the program and compare.

Show the answer

The output is:

*   *   *
 * * * * 
  *   *  

Row 0 has stars in columns 0, 4 and 8. Row 1 has stars in columns 1, 3, 5 and 7. Row 2 has stars in columns 2 and 6.

Exercise 3: a row function

Write a function row_at(step, rows) that returns the row for a given step, without a loop. Use the cycle idea from row_of. Check it for steps 0 to 7 with rows set to 3.

Show the answer

def row_at(step, rows):
    cycle = 2 * (rows - 1)
    position = step % cycle
    return position if position < rows else cycle - position
 
 
print([row_at(step, 3) for step in range(8)])

The output is:

[0, 1, 2, 1, 0, 1, 2, 1]

This is the same rule as row_of in Example 3. The position if position < rows else cycle - position line is a conditional expression. It picks one of two values in a single line.

Exercise 4: predict the zigzag

Work out convert("ABCD", 2) on paper before you run it. Write the letter that goes in each row, then join the rows.

Show the answer

A goes in row 0, B in row 1, C in row 0 and D in row 1. Row 0 holds AC and row 1 holds BD. The result is ACBD.

Exercise 5: fix the one-row crash

The convert function in Common mistakes fails for convert("AB", 1). Add the fix and state what the function returns.

Show the answer

Add this line at the top of convert:

if num_rows == 1:
    return text

The function returns AB. The early return skips the loop over the letters, so the marker never moves to row 1.

Summary

• A zigzag is a position that moves one step at a time and turns around at each edge.
• Keep the direction in a variable with the value 1 or -1. Check the edges first, then add the direction to the row.
• Set the direction to a fixed value at each edge. Reversing it is easy to get wrong.
• A look-ahead check computes the next row and turns around if that row is outside 0 to rows - 1. It works for any starting row.
• Python accepts negative indexes, so a wrong row number can pick a real row without an error. Check the edges.
• To write text in a zigzag, append each letter to the list of its row, then join the rows top to bottom.
• A one-row board needs a special case, because the marker cannot move.