Solving a numbers puzzle with Python
Recently, at Airteam, we had a contest that was inspired by a show called Letters and Numbers.
To give you an idea of what it is, watch this video:
Our contest had slightly different rules (without any concept of limiting “large” or “small” numbers), and basically goes like this:
Rules of the game
Given the numbers 7, 2, 9, and 7, find the best way to use these numbers to arrive closest to the number 211. You don’t have to use all numbers, but you may not repeat numbers.
You may use any operator you wish, but not ones that use special constants, like the exponential constant. If you use power, e.g 7^2, both 7 and 2 are considered numbers.
Being horrible at mental (and handwritten) math, I decided to use Python to solve this problem.
(For the record, I did not win, but came close)
I thought it was interesting enough to talk about.
The problem in more detail
Let’s talk about the data we have, and the operations that are allowed.
We’re given, essentially, a list of numbers. In this case it’s [7, 2, 9, 7]
.
We also have a set of allowed operators.
Most operators apply to two numbers. For example add
is an operator:
from operator import add
add(2, 2) # returns 4
In other words, I need to find the correct combination of operators and numbers to arrive at the target value, which is 211.
Modelling the problem
So, how do we begin to solve this in programming terms?
If you know some LISP (or dialects like Clojure), you know that we can express an operator as follows:
(add 1 2) ;this represents `1 + 2`
And we can also build an arbitrary expression by nesting these structures:
(mul 3 (add 1 2)) ;this represents 3 * (1 + 2)
As you can see, we can represent any mathematical expression using nested lists.
This means that one of these nested lists is a candidate for a solution. Once evaluated we can determine what value this candidate resolves to.
That means we also have to write function to evaluate the candidates.
The strategy, broadly, is:
- Generate an exhaustive list of candidates from the given numbers
- Go through each of these candidates and parse them.
- Return the best solution (if an exact one is not found)
I decided to go further, and try to return every possible exact solution, if one or more are found.
Let’s break this down step-by-step.
Permutations
Thanks to Python’s itertools
module, generating all permutations of a list is relatively straightforward:
import itertools
itertools.permutations([1, 2, 3], 3)
In other words, I can generate all possible orderings of the given list.
Given any random ordering, I can then iterate over all operations and build all possible candidates.
Grouping numbers together
Merely generating permutations is unfortunately not enough.
Recall that an operator is applied to two numbers.
That means there exists some valid combinations of groupings, for each permutation.
For a single permutation of four numbers [a, b, c, d]
, here are some of the possible groupings:
(a, (b, (c, d)))
((a, b), (c, d))
(((a, b), c), d)
((a, (b, c)), d)
Here’s the code I used to generate these groupings:
def _get_groupings(nums):
if len(nums) == 1:
yield nums[0]
elif len(nums) == 2:
yield nums
else:
yield from (
(xx, yy)
for i in range(1, len(nums))
for xx in _get_groupings(nums[:i])
for yy in _get_groupings(nums[i:])
)
def get_groupings(perms):
for nums in perms:
yield from _get_groupings(nums)
The function get_groupings
receives the exhaustive list of permutations, and for each iteration we generate the groupings.
This is a recursive function, as is most of the functions written for this task.
The way this works is it treats a single permutation as two parts, let’s call them x
and y
.
Either x
or y
could be a single value, or a tuple. The terminating case is if we find a valid “unit”, which is either a group of 2 numbers, or just a single number.
Generating candidates
Now that I have a function to generate groupings from permutations, I can now write a function to generate the permutations, collect the groupings, and transform them into a valid expression.
This is the code I arrived at:
from itertools import permutations
OPERATIONS = ['add', 'sub', 'mul', 'div']
def _generate_candidates(nums):
x, y = nums[0], nums[1]
if not isinstance(x, tuple) and not isinstance(y, tuple):
for name in OPERATIONS:
yield (name, x, y)
else:
x_gens = [x] if not isinstance(x, tuple) else _generate_candidates(x)
y_gens = [y] if not isinstance(y, tuple) else _generate_candidates(y)
yield from (
(name, a, b) for a in x_gens for b in y_gens for name in OPERATIONS
)
def generate_candidates(numbers):
all_permutations = (
x for r in range(2, len(numbers) + 1) for x in permutations(numbers, r)
)
for g in get_groupings(all_permutations):
yield from _generate_candidates(g)
There’s quite a bit is going on there.
Recall that we don’t need to use all numbers. So we also need permutations of length 2 up to N where N is the length of the list.
That’s why I wrote this:
all_permutations = (
x for r in range(2, len(numbers) + 1) for x in permutations(numbers, r)
)
The variable all_permutations
is now a generator object which will yield all possible permutations at all possible lengths.
We can then pass this to the get_groupings
to receive the valid groupings for each permutation.
At this point we are still only dealing with tuples of numbers, e.g (1, (2, 3))
. Remember we still want something of the form ('mul, 3, ('add', 1, 2))
.
That’s where the private function _generate_candidates()
comes in.
Let’s zoom in at this function:
def _generate_candidates(nums):
x, y = nums[0], nums[1]
if not isinstance(x, tuple) and not isinstance(y, tuple):
for name in OPERATIONS:
yield (name, x, y)
else:
x_gens = [x] if not isinstance(x, tuple) else _generate_candidates(x)
y_gens = [y] if not isinstance(y, tuple) else _generate_candidates(y)
yield from (
(name, a, b) for a in x_gens for b in y_gens for name in OPERATIONS
)
This is another recursive function.
Essentially, this tries to find the smallest grouping unit we can apply an operator to, and tries to build up a more complex expression.
Unfortunately, there was a bit of type checking in here which hurt readability. If I were to implement this again, I would have written a class.
Evaluating candidates
We still need some mechanism of evaluating the value of a something like ('mul', ('add', 1, 2), 3)
.
To do this, I wrote another recursive solution:
import functools
import operator
OPERATIONS = {
'add': operator.add,
'sub': operator.sub,
'mul': operator.mul,
'div': operator.truediv,
}
@functools.lru_cache()
def compute(candidate):
name, x, y = candidate
child_x, child_y = x, y
child_x = compute(x)[1] if isinstance(x, tuple) else x
child_y = compute(y)[1] if isinstance(y, tuple) else y
return candidate, OPERATIONS[name](child_x, child_y)
def get_best_candidates(candidates, final_result):
best_candidate, best_result, exact_found = None, 0, False
for candidate in candidates:
try:
n, r = compute(candidate)
except (ValueError, ZeroDivisionError, TypeError, OverflowError):
continue
if abs(r - final_result) < abs(best_result - final_result):
best_candidate, best_result = n, r
if r == final_result:
exact_found = True
yield r, n
if not exact_found:
yield best_result, best_candidate
There are two functions, get_best_candidates()
and compute()
.
The function get_best_candidates()
is just boring procedural code, and isn’t too interesting, it just tries to find the best solution from the candidates. If something happens (like a ZeroDivisionError
), we simply ignore the candidate and move on.
The private function compute()
is the more interesting recursive one. Broadly, we can think of a candidate expression like this:
(OPERATOR, val1, val2)
The values, val1
and val2
can be numbers, or they can be nested expressions.
As mentioned, recursive functions have terminating conditions. The terminating condition is when both values are numbers. If they are both numbers, we can evaluate this directly.
Notice at the top we have this piece of code:
OPERATIONS = {
'add': operator.add,
'sub': operator.sub,
'mul': operator.mul,
'div': operator.truediv,
}
This is dict
where the keys are the names of the operators (this is the first element in each tuple), and the values are functions. That means we can do this:
OPERATIONS['add'](1, 2) # this does 1 + 2
If, instead, one (or both) of the values in the expression is another expression, we can recursively call compute()
on the inner expression.
Also, you may notice that I used @functools.lru_cache()
to decorate compute()
. This caches any calculations that get reused over and over again.
So that’s the whole algorithm!
Additional operators
In the code above, I only included basic arithmetic.
Actually, there was nothing in the rules that forbade using integer or bitwise operators.
So, I naively decided to add some more operators:
OPERATIONS = {
'add': operator.add,
'sub': operator.sub,
'mul': operator.mul,
'div': operator.truediv,
'floordiv': operator.floordiv,
'pow': operator.pow,
'mod': operator.mod,
'lshift': operator.lshift,
'rshift': operator.rshift,
}
That was pretty disastrous.
Some of the values can get pretty big. So sometimes we see cases where things like this happen:
2 ** 123445345345
7 << 123123123355
This basically killed the process.
To get around this, I had to write custom operators:
def hacked_pow(a, b):
if b > MAX_POWER:
raise ValueError("let's not use such high powers")
return operator.pow(a, b)
def hacked_lshift(a, b):
if b > MAX_LSHIFT:
raise ValueError("let's not shift this much")
return operator.lshift(a, b)
def hacked_rshift(a, b):
if b > MAX_RSHIFT:
raise ValueError("let's not shift this much")
return operator.rshift(a, b)
This worked a lot better.
Parsing the expression
I can get the solution from reading the tuple that represents the expression.
However, it’s in this form: ('mul, ('pow', ('add', 1, 2), 3), 7)
That’s pretty hard to read, and I only have 4 numbers in that example.
What I want instead is something that displays it like this: (1 + 2) ^ 3 * 7
So I wrote a parser:
SYMBOLS = {
"add": "+",
"mul": "*",
"sub": "-",
"div": "/",
"floordiv": "//",
"lshift": "<<",
"rshift": ">>",
"mod": "%",
"pow": "^",
}
def parse(formula, bracket=True):
name, x, y = formula
symbol = SYMBOLS[name]
child_x = parse(x) if isinstance(x, tuple) else str(x)
child_y = parse(y) if isinstance(y, tuple) else str(y)
formula = f" {symbol} ".join([child_x, child_y])
if bracket:
return f"({formula})"
return formula
Putting it together, I tried the values given in the video above:
input_numbers = (75, 25, 50, 100, 8, 2)
target = 431
for (x, y) in compute(generate_candidates(input_numbers), target):
expr = parse(y, False)
print(f"{expr} = {x}")
And here’s the output:
75 + (100 + (2 ^ 8)) = 431
(75 + 100) + (2 ^ 8) = 431
75 + ((2 ^ 8) + 100) = 431
(75 + (2 ^ 8)) + 100 = 431
100 + (75 + (2 ^ 8)) = 431
(100 + 75) + (2 ^ 8) = 431
100 + ((2 ^ 8) + 75) = 431
(100 + (2 ^ 8)) + 75 = 431
((2 ^ 8) + 75) + 100 = 431
((2 ^ 8) + 100) + 75 = 431
(((75 * 25) - 50) - 100) >> 2 = 431
75 // (25 ^ (50 / (8 - 100))) = 431.0
(75 * (25 - (50 % 8))) >> 2 = 431
((75 + (25 * 50)) >> 2) + 100 = 431
(75 * (25 - (100 // 50))) >> 2 = 431
(((75 * 25) - 100) - 50) >> 2 = 431
(75 + ((25 + 8) * 50)) >> 2 = 431
(75 * ((25 ^ 2) - 50)) // 100 = 431
...
And it goes on.
Pretty nice!
Conclusion and thoughts
This was a really fun project! I didn’t have the time to worry about code quality, and as such the code you see here isn’t the most readable. If I were to re-write this, I might be writing more classes instead of type checking for tuples.
If any beginners are reading and wonder why recursion is useful in Python, this was a really great example of how it can come useful.
Note that you don’t have to use recursion for this problem (you could use a stack), but that would make things a little more tricky to manage, in my opinion.
Also, this is basically a brute force solution, which is probably not very good. My search space grows O(n!)
(due to generating permutations).
I think there are better strategies to solve this problem that I have yet to explore. For example, we could do things more intelligently by starting out with the target number, and then moving backwards by performing operations on the available numbers.
My code is available here, if you want to take a look at it: https://github.com/yoongkang/numbers_puzzle