Count number of ways to paint a fence with N posts using K colors

A friend sent me this question, and I'd love somebody to review this and give his opinion.


Write an algorithm that counts the number of ways you can paint a
fence with N posts using K colors such that no more than 2 adjacent
fence posts are painted with the same color"

Is this algorithm accurate? Could it be implemented in a better way? Is there anything really un-pythonic that I did there? I hope I haven't made any embarrassing mistakes.

def find_combinations(n, colours):

combos_list = []

    for i in range(1, n+1):

        #Check if the combos_list is empty (happens when post number == 1)
        if not combos_list:
            #Add all variations for when there's only one post
            for c in colours:
                combos_list.append([c])
        else:
            new_combo_list = []
            for combo in combos_list:
                for c in colours:
                    #Check if last 2 posts are from the same colour, meaning we can't add another one
                    if not combo[len(combo) - 2:] == [c, c]:
                        combo_to_add = combo[:]
                        combo_to_add.append(c)
                        new_combo_list.append(combo_to_add)
            combos_list = new_combo_list[:]

    return combos_list, len(combos_list)

colours = ["black", "white"]
n = 4

print find_combinations(n, colours)

Example of input:


colours = ["a", "b"]

n = 3

Output:

• a, a, b

• a, b, a

• a, b, b

• b, a, a

• b, a, b

• b, b, a

Number of combos - 6

1. Code review

-
The question says, "Write an algorithm that counts the number of ways" but the program does not only do this: it also builds a list of the combinations. This means that the program is limited to values of \$N\$ and \$K\$ that are small enough that the list of combinations fits in the storage space of the computer. Whereas a program that did not have to produce this list could work for much larger values of \$N\$ and \$K\$, as we'll see in §3 below.

-
It's good practice to write a docstring for each function, describing what the function does, the arguments that it takes, and what it returns.

-
If a function returns a list, it does not also need to return the length of the list. If the caller needs that information, it can call len just as easily. So the function could just return combos_list.

-
In the line:

combos_list = new_combo_list[:]

it's unnecessary (and wasteful) to take a copy of new_combo_list. The new list can just be assigned to the old name, like this:

combos_list = new_combo_list

-
When a loop index (here i) is not used, it's conventional to name it _.

-
The code that's guarded by the line:

if not combos_list:

only runs once (on the first iteration of the loop), so it would be simpler to move it out of the loop, have one fewer iteration of the loop, and so avoid the if ... else .... Like this:

combos_list = []
for c in colours:
    combos_list.append([c])
for _ in range(1, n):
    new_combo_list = []
    # etc.

-
Instead of:

combos_list = []
for c in colours:
    combos_list.append([c])

use a list comprehension and write:

combos_list = [[c] for c in colours]

-
These lines:

combo_to_add = combo[:]
combo_to_add.append(c)
new_combo_list.append(combo_to_add)

could be simplified to:

new_combo_list.append(combo + [c])

-
This condition:

if not combo[len(combo) - 2:] == [c, c]:

could be simplified to:

if combo[-2:] != [c, c]:

(since Python allows you index a list backwards from the end using negative numbers).

1. Revised code

def find_combinations(n, colours):
    """Return a list of combinations of n items from the sequence colours,
    with the constraint that no colour appears three times in a row.

    >>> [''.join(combo) for combo in find_combinations(3, 'ab')]
    ['aab', 'aba', 'abb', 'baa', 'bab', 'bba']

    """
    combos = [[c] for c in colours]

    for _ in range(1, n):
        new_combos = []
        for combo in combos:
            for c in colours:
                # Check that last two colours aren't both c.
                if combo[-2:] != [c, c]:
                    new_combos.append(combo + [c])
        combos = new_combos

    return combos

Note the example code in the docstring. This can be run as a test case using the doctest module.

1. Let me count the ways

The question asked for "an algorithm that counts the number of ways", so here's how I'd go about that.

I'll start with some notation. Assume that \$k\$ is fixed. Then:

-
\$T(n)\$ is the total number of ways to paint \$n\$ fenceposts with \$k\$ colours, subject to the constraint that no three adjacent posts are the same colour.

-
\$D(n)\$ is the number of ways to paint \$n\$ fenceposts with \$k\$ colours such that the last two fenceposts are different colours (or \$n

-
\$S(n)\$ is the number of ways to paint \$n\$ fenceposts with \$k\$ colours such that the last two fenceposts are the same colour.

And then the following recurrence relations are clear:

• \$D(n) = (k−1)(D(n−1) + S(n−1))\$ for \$n ≥ 2\$

• \$S(n) = D(n−1)\$ for \$n ≥ 2\$

• \$T(n) = D(n) + S(n)\$

We can use (2) to eliminate \$S\$ (being careful about the bound on \$n\$ as we do so), getting the recurrences:

• \$D(n) = (k−1)(D(n−1) + D(n−2))\$ for \$n ≥ 3\$

• \$T(n) = D(n) + D(n−1)\$ for \$n ≥ 2\$

And then substitute (5) into (4), getting:

• \$D(n) = (k−1)(T(n−1))\$ for \$n ≥ 3\$

and so:

• \$T(n) = (k−1)(T(n−1) + T(n−2))\$ for \$n ≥ 3\$

Since the recurrence depends only on the last two values of the sequence, it's easy to program:

def triple_free_combinations(n, k):
    """Return the number of ways to choose n items (with k choices for
    each item), subject to the constraint that no colour appears three
    times in a row.

    """
    if n == 0:
        return 1
    a, b = k, k * k
    for _ in range(n - 1):
        a, b = b, (k - 1) * (a + b)
    return a

(You should recognize this as a very simple example of the tabular method, aka "dynamic programming".)

The effort that went into find_combinations is not wasted, because now we can use that function as source of test cases:

>>> from itertools import product
>>> all(triple_free_combinations(n, k)
... == len(find_combinations(n, list(range(k))))
... for n, k in product(range(1, 8), repeat=2))
True


Another useful source of test cases is the On-Line Encyclopedia of Integer Sequences. For example, when \$k=2\$ we hav

StackExchange Code Review Q#63614, answer score: 20