Connect Four implemented in 3 lines of Python

A tiny implementation of Connect 4 on the terminal, with explanation.

by on

A little while ago, I wrote an implementation of Hangman in 3 lines of Python and had a lot of fun working on it. I was surprised to find that other hackers on the Internet also enjoyed reading, playing with, and riffing on it. So, I’ve decided to perform the same trick, but this time with a (slightly) harder game: Connect Four.

How Connect Four is Played


Fig. 1 A connect four board (board position not possible during normal play)

Connect Four is a popular game played on a grid that is 7 across, and 6 deep. Players take turns inserting pieces into the grid, with the aim of being the first to connect four pieces of their own colour. Unlike Tic-Tac-Toe, however, players do not have the freedom to place their piece in any free cell they wish. Instead, players select a given column, and the piece must occupy the lowest free cell within that column.

On most physical models of the Connect Four game, the grid is positioned vertically so that gravity pulls pieces down the columns into the lowest spot they can occupy.

One other fact about Connect Four that I should mention : Connect Four has been solved, in much the same way that Tic-Tac-Toe has been solved. And, unlike Tic-Tac-Toe, where perfect play from both players results in a draw, in Connect Four the first player to move can always force a win. This makes it a boring game to play against a computer using the perfect algorithm but against other humans it can still be fun.

The solution itself

Here is the ugly, ugly solution. You may need to scroll a bit to see it all.

licence, itertools, board, vectors, get, put, coloured, board_string, winner = 'MIT', __import__('itertools'), [[] for _ in xrange(7)], [(0, 1), (1, 0), (1, 1), (1, -1)], lambda i, j: board[i][j] if 0 <= i < 7 and len(board[i]) > j >= 0 else ' ', lambda i, p: board[i].append(p) or True if 0 <= i < 7 and len(board[i]) < 6 else False, (lambda string, colour: {'red': '\033[1;31m{}\033[1;m', 'yellow': '\033[1;33m{}\033[1;m', 'blue': '\033[1;34m{}\033[1;m'}.get(colour, '{}').format(string)) if not __import__('os').getenv('NOCOLOR') else lambda string, colour: string, lambda: '\n'.join(coloured(' | ', 'blue') + coloured(' | ', 'blue').join(get(column, row) for column in range(7)) + coloured(' | ', 'blue') for row in range(5, -1, -1)) + '\n' + coloured(' | ', 'blue') + coloured(' | ', 'blue').join(coloured('{}', 'blue').format(i + 1) for i in range(7)) + coloured(' | ', 'blue'), lambda: any(all(get(i, j) != ' ' and get(i, j) == get(i + vi * step, j + vj * step) for step in range(4)) for vi, vj in vectors for j in range(6) for i in range(7))
for player, players in (lambda players: itertools.izip(itertools.takewhile(lambda _: not winner() and not all(len(col) == 6 for col in board), players), (players for _ in itertools.count())))(itertools.cycle([coloured('X', 'red'), coloured('O', 'yellow')])): _ if put(int(''.join(filter(str.isdigit, raw_input(board_string() + '\nPlayer %s, pick a column to drop into: ' % player))) or -1) - 1, player) else players.next()
print board_string(), " ".join(["\n", player, "wins!"]) if winner() else "\nDraw!"

Once again, those three lines of Python 2.7 are all you need to implement a non-trivial game. If you want to run it at your own terminal to play it yourself, you can download it from this gist, or copy and paste the 3 lines above into connect4.py and run it. If your terminal doesn’t support ANSI colour escape sequences you might see some garbage output that looks like

\033[1;33m|\033[1;m\033[1;33m|\033[1;m\033[1;33m|\033[1;m\033[1;...

In that case, you can enable no-color mode by passing in an environment variable like this:

NOCOLOR=True python connect4.py

Here’s all the features I was able to pack into those 3 lines:

  • Printing a representation of the board state on every turn
  • An optional colour mode
  • Keeping track of the active player
  • Rejecting erroneous inputs, such as "potato" (I decided that "potato1" should be interpreted as column 1)
  • Rejecting invalid inputs, like -1 or trying to add a piece to a column that was already filled
  • Maintaining the same active player if an invalid move is made
  • Win condition detection
  • Draw detection (but only after the entire board has been filled)

So how did I do it? Stay tuned for a line-by-line analysis.

Why 3 lines?

But before we dive into the analysis, I’d like to explain why I find 3 lines of Python a satisfying target. Humans are naturally inclined to think in threes. Three bears, three little pigs, three princes off to seek their fortune. Good stories have a beginning, a middle, and an end1.

In the same way, 3 lines of Python break down in a really clear way the different parts of my program:

  • a giant assignment, which takes care of imports, sets up the board state and defines some helper functions
  • the main game loop, which evaluates the board, asks for input and changes the state of the game
  • the code to congratulate the winner, which runs when the game is finished

Of course, any program of the form:

a = 10
while a > 2: a -= 1
print a

can be turned into a single line in the following way:

print (lambda fn, arg: fn(fn, a))(
    lambda fn, a: fn(a - 1, fn) if a > 2 else a, 10)

I won’t go into too much detail about how that magic works in this post, but I am aware it can be done. Still, I’m going to leave my Connect Four implementation at 3 lines, and leave collapsing it into one as an exercise for enthusiastic readers.

Analysis

The three lines of the program separately correspond to the setup, the main game loop, and the final winner declaration of the game. However, compared to my previous game of hangman, there’s rather a bit more going on within each line. To better explain the program, I’m going to run through an overview of each line, and then explain the logical parts together.

The first line just uses Python’s multiple assignment to state the license, import required dependencies, initialize the game state, and define some helper functions. As you can see, using lambda expressions are a necessity because function definitions would otherwise waste a line. Don’t worry too much about the content of the functions for now, as we’ll examine them in more detail later.

Here’s the first line again:

licence, itertools, board, vectors, get, put, coloured, board_string = 'MIT', __import__('itertools'), [[] for _ in xrange(7)], [(0, 1), (1, 0), (1, 1), (1, -1)], lambda i, j: board[i][j] if 0 <= i < 7 and len(board[i]) > j >= 0 else ' ', lambda i, p: board[i].append(p) or True if 0 <= i < 7 and len(board[i]) < 6 else False, lambda string, colour: {'red': '\033[1;31m{}\033[1;m', 'yellow': '\033[1;33m{}\033[1;m', 'blue': '\033[1;34m{}\033[1;m'}.get(colour, '{}').format(string) if not __import__('os').getenv('NOCOLOR') else string, lambda: '\n'.join(coloured(' | ', 'blue') + coloured(' | ', 'blue').join(get(column, row) for column in range(7)) + coloured(' | ', 'blue') for row in range(5, -1, -1)) + '\n' + coloured(' | ', 'blue') + coloured(' | ', 'blue').join(coloured('{}', 'blue').format(i + 1) for i in range(7)) + coloured(' | ', 'blue')

And here it is rendered slightly more intelligible by using each assignment on one line, and by adding line breaks where needed for clarity:

licence = 'MIT'
itertools = __import__('itertools')
board = [[] for _ in xrange(7)]
vectors = [(0, 1), (1, 0), (1, 1), (1, -1)]
get = lambda i, j: (
    board[i][j] if
    0 <= i < 7 and len(board[i]) > j >= 0
    else ' ' )
put = lambda i, p: (
    board[i].append(p) or True if
    0 <= i < 7 and len(board[i]) < 6
    else False )
coloured = ((lambda string, colour: {
    'red': '\033[1;31m{}\033[1;m',
    'yellow': '\033[1;33m{}\033[1;m',
    'blue': '\033[1;34m{}\033[1;m'}.get(colour, '{}').format(string)) if
    not __import__('os').getenv('NOCOLOR')
    else lambda string, colour: string)
board_string = (lambda:
    '\n'.join(
        coloured(' | ', 'blue') +
        coloured(' | ', 'blue').join(
            get(board, column, row) for column in range(7)) +
        coloured(' | ', 'blue')
        for row in range(5, -1, -1)) +
    '\n' +
    coloured(' | ', 'blue') +
    coloured(' | ', 'blue').join(
        coloured('{}', 'blue').format(i + 1) for i in range(7)) +
    coloured(' | ', 'blue'))
winner = lambda: any(
    all(get(i, j) != ' ' and
        get(i, j) == get(i + vi * step, j + vj * step)
        for step in range(4))
    for vi, vj in vectors
    for j in range(6)
    for i in range(7))

The one tricky bit I want to point out right now is the use of the builtin __import__ function to get away with importing in an expression, instead of having to waste a line on an import statement.

The second line is a giant while loop, even though it uses the for keyword. Python allows (but frowns upon) collapsing the body of a single statement for or while loop onto one line, which I take advantage of in this program.

for player, players in (lambda players: itertools.izip(itertools.takewhile(lambda _: not any(all(get(i, j) != ' ' and get(i, j) == get(i + vi * step, j + vj * step) for step in range(4)) for vi, vj in vectors for j in range(6) for i in range(7)) and not all(len(col) == 6 for col in board), players), (players for _ in itertools.count())))(itertools.cycle([coloured('X', 'red'), coloured('O', 'yellow')])): _ if put(int(''.join(filter(str.isdigit, raw_input(board_string() + '\nPlayer %s, pick a column to drop into: ' % player))) or -1) - 1, player) else players.next()

We’ll revisit this line in greater detail later, but for now this translation of its overall structure is sufficient:

players = itertools.cycle(['X', 'O'])
while not winner() and not have_draw():
      # The current player is the next one in rotation.
      current_player = players.next()
      print board_string()
      column = read_input()
      if put(column):
          pass
      else:
          # Inserting failed, give current_player a
          # turn again by skipping the next player's turn.
          players.next()

The final line is very simple:

print board_string(), " ".join(["\n", player, "wins!"]) if winner() else "\nDraw!"

It just prints the board a final time, and then announces the winner if there is one, or calls a draw.

The game board

The game board is initialized as a list of empty lists. Each inner empty list represents a column position. I chose this representation over a full 7x6 2-D array because it lets me append to the end of the current list at each column as players keep adding pieces to the board, without having to perform any extra work to store or calculate the index of the highest free cell in every column. Instead, the function put() can simply index into the board and append to a given column.

Written more idiomatically, put() looks like this:

def put(column, piece):
    if 0 <= i < 7 and len(board[i]) < 6:
        b[i].append(p)
        return True
    else:
        return False

Put returns a boolean value depending on if it was successful at adding a piece or not. If the user enters a bogus value, or attempts to play in a column that is already at capacity, this function does nothing and returns False. The program uses this to detect that input could not be read from the current player, and gives that player another turn.

The body of a lambda expression in Python can only be one expression. Luckily for us, Python allows us to express a branch of the form

def foo(a):
    if a:
        return 2 + 2
    else:
        return a ** 2

with a single ternary expression as

def foo(a):
    return 2 + 2 if a else a ** 2

However, each element in this expression must itself be an expression, and we need to do two things if the append is safe to go ahead: actually do the append, and then return a truthy value. So we abuse Python’s left-to-right logic evaluation to do both, since append() returns None, and None or True always evaluates to True.

return board[i].append(piece) or True

get() is very similar to put(). It’s defined to take two parameters, the column and the row index. Just as put() attempts to be safe in the face of improper input, get() always returns the safe value of an empty space, " ", if you ask it to retrieve a value that is out of bounds.

board_string(), as you might expect, returns the string representation of the board. We defined this function in the first line, along with all the other initialization. Even though the function board_string is defined, in some sense, simultaneously, with the variables it depends on, viz. board, get and coloured, due to Python’s late binding board_string is able to resolve them correctly when it is run later. This is a great bonus to brevity because it lets us implicitly shared state between the various board functions—but an overall impediment to maintainability.

board_string() simply builds a string by iterating over every cell in the board, and joining the results together into the right format. It takes advantage of the sane behaviour of get() when called for cell indexes that are out of bounds, so that it doesn’t have to duplicate bounds checking. Relying on this behaviour from get() leads to an even greater simplification in win detection.

Win detection

This simple function returns the character of the winning player, if there is one, or None if there isn’t.

winner = lambda: any(
    all(get(i, j) != ' ' and
        get(i, j) == get(i + vi * step, j + vj * step)
        for step in range(4))
    for vi, vj in vectors
    for j in range(6) for i in range(7))

There are four vectors along which a winning play might be made.

Vertically
Horizontally
Diagonally upward
Diagonally downward

The win detection function is a nested for-loop that checks every single segment of four cells beginning at every single cell on the board–including “virtual segments” that go off the board.

for vi, vj in vectors:
    for j in range(6):
        for i in range(7):
            # Build up a list of all the cells to
            # be checked in this iteration.
            cells = []
            for step in range(4):
                cells.append(get(i + vi * step, j + vj * step))
            # If these four cells are all the same, and all
            # non-blank, we have a winner--return it!
            if all(cell != ' ' and cell == cells[0] for cell in cells):
                return cell[0]

Translating the winner() function out by hand like this somewhat overstates how much extra work it performs. In the one-line invocation, all gets directly applied to the generator for the actual cells, and thanks to the fact that it shortcuts logic when possible, it never goes on seeking past the first empty or non-matching piece.

However, I have to admit that in designing this win function, I was aiming for brevity over performance. On every invocation, this function performs an exhaustive search of the board and a little bit more. A more performance-conscious definition could rely on the constraints that the only winner possible is the player who last moved, and that any winning arrangement created must go through the last played location, to greatly cut down on the search space.

The main loop

The code of the main game loop looks like this:

for player, players in (lambda players: itertools.izip(itertools.takewhile(lambda _: not winner() and not all(len(col) == 6 for col in board), players), (players for _ in itertools.count())))(itertools.cycle([coloured('X', 'red'), coloured('O', 'yellow')])): _ if put(int(''.join(filter(str.isdigit, raw_input(board_string() + '\nPlayer %s, pick a column to drop into: ' % player))) or -1) - 1, player) else players.next()

Written more legibly, it would look something like:

players = itertools.cycle(['X', 'O'])
while not winner() and not have_draw():
      # The current player is the next one in rotation.
      current_player = players.next()
      print board_string()
      column = read_input()
      if put(column):
          pass
      else:
          # Inserting failed, give current_player a turn again
          # by skipping the next player's turn.
          players.next()

players is an infinite cycle of ['X', 'O', 'X', 'O', ...]. On every iteration of this loop, we keep pulling the next player off the list, and offering them a chance to make a move. If their move is invalid, we skip the other player so they can get a chance to enter a correct move.

There is an extra bit of complexity that the one-liner has to account for. The inner body of the loop needs to have access to the players iterator, so that it can give the current player a turn again if they make the wrong move.

The function itertools.takewhile will iterate over the players iterator as long as the game is still going, but returns only the current player–access to the players iterator will be “lost” if the only reference is passed into takewhile. So, I create a construct that will both handle termination correctly, and give me access to to the underlying players iterator.

That’s why this unwieldy line is needed:

(lambda players: itertools.izip(
    itertools.takewhile(lambda _: not winner() and not all(len(col) == 6 for col in board),
                        players),
    (players for _ in itertools.count()))
    )(itertools.cycle(['X', 'O']))

It creates a lambda function that takes an infinite cycle of players, and returns a generator that will yield a tuple containing the next player as well as a reference to that infinite cycle. That function is then called with our infinite player cycle, to create such a generator. That generator will stop yielding tuples as soon as the board is either won, or filled.

Colouring

The coloured() function takes care of colourizing the output. There are actually two versions of this function, depending on whether colour output is required or not. If the environment variable NOCOLOR2 is set, the following check runs at initialization time

coloured = (<coloured-implementation>
            if not __import__('os').getenv('NOCOLOR')
            else lambda string, colour: string)

and the function is replaced by one that does nothing. Otherwise, the function is defined to simply look up the ANSI Colour Escape code for one of the three defined colours, and format the given string according to it:

(lambda string, colour: {
    'red': '\033[1;31m{}\033[1;m',
    'yellow': '\033[1;33m{}\033[1;m',
    'blue': '\033[1;34m{}\033[1;m' }.get(colour, '{}').format(string))

Since we don’t expect the environment variable to change during the execution of the program, it’s more efficient to perform this check once and swap functions, rather than checking whether the environment variable is set every time within the body of the function.

Conclusion

It wasn’t much work to create a fully playable connect four in 3 lines of Python. However, as with Hangman, explaining exactly what’s going on within this program took many more lines than that. In fact, although I would not consider Connect Four to be much more complicated than Hangman, explaining it took disproportionately more effort.

I would guess that the amount of explanation required is not just proportional to the complexity of the program, but actually scales according to the “compression-factor”, as it were. This exercise in extreme terseness might teach us a lesson for normal everyday programming—that programs are at their most readable when they are written at the appropriate level of terseness, a level that matches the complexity of the problem.

Finally, as I was writing this, I kept thinking that it would be simple to create an automatic tool, if none exist already, that could take an arbitrary Python program and compress it down to a single line. However, just as the fact that Connect Four has been “solved” does not stop it from being played for fun, the existence of tools to minify Python doesn’t rob me of the enjoyment I received from this process.


  1. Granted, the stories of today are also entitled to include an additional end-credits scene with a sequel hook. ↩︎
  2. You might be wonding why I spell the word colour differently in the program. The variable in Python with the -our suffix, colour, and the environment variable with -OR, COLOR? I’ve been raised to always spell the word with the -our suffix, but I know that U.S. audiences (and computer variables in general), tend to spell it -OR. My compromise is that anything I write for myself (and variables in nearly-obfuscated code counts as this) I spell -our, but anything I write that needs to be predictable, be it production code or environment variables that my users will have to edit, I spell the American Way. ↩︎

Other articles you may like

This article was filed under:
Programming Python Code Golf