Countdown numbers game (Solution generator)

Question

I am trying to program a solution generator for the Countdown numbers game in Python, following some guidance from DataGenetics -- those unfamiliar with the rules, you can review them here as well -- but I wrote the code myself.

Code

def postfix_to_infix(postfix):
 """
 :param postfix: tuple containing solution in reverse polish notation
 :return: string of solution in polish notation
 """
 soln = list(postfix)
 while True:
 try:
 op_idx, op = [(idx, tkn) for idx, tkn in enumerate(soln) if tkn in ['+', '-', '*', '/']][0]
 except IndexError:
 break
 else:
 block = '(' + str(soln[op_idx - 2]) + op + str(soln[op_idx - 1]) + ')'
 soln[op_idx - 2:op_idx + 1] = [''.join(block)]
 return soln
def main():
 import random
 # Print welcome message
 print("""
 =====================================
 Welcome to Countdown, Python version!
 =====================================
 """)
 # *** <Start game> ***
 play_new_game = True
 while play_new_game:
 # Generate the set of possible numbers allowed by Countdown rules
 large_numbers = [25, 50, 75, 100]
 small_numbers = list(range(1, 11)) * 2
 # Generate six number tiles according to user input for large numbers
 while True:
 try:
 usr_large_num = int(input("How many LARGE numbers do you want (0-4)? "))
 except ValueError:
 pass
 else:
 if 0 <= usr_large_num <= 4:
 break
 six_tiles = random.sample(large_numbers, usr_large_num) + \
 random.sample(small_numbers, 6 - usr_large_num)
 # Generate random target output: any three-digit number
 target = random.randint(100, 999)
 print(six_tiles)
 print(f"Target: {target}")
 # Algorithm to find solution to target
 import itertools as it
 from more_itertools import distinct_permutations
 import operator
 ops_tiles = {"+": operator.add,
 "-": operator.sub,
 "*": operator.mul,
 "/": operator.truediv}
 for nums in distinct_permutations(six_tiles):
 for ops in it.combinations_with_replacement(ops_tiles, 4):
 # Generate the list of postfix's (6 numbers + operators)
 # by permuting over all elements,
 # But removing those postfix's that begin with two operators
 postfix_gen = (pf for pf in distinct_permutations(nums + ops) if
 pf[0] not in ops_tiles and pf[1] not in ops_tiles)
 for postfix in postfix_gen:
 # Implement postfix algorithm
 stack = list(postfix[0:2])
 for idx, token in enumerate(postfix[2:], start=2):
 # If token is an operator...
 if token in ops_tiles:
 # Take the last 2 tokens in the stack
 try:
 operand_2 = stack.pop()
 operand_1 = stack.pop()
 except IndexError:
 break
 # And operate on the 2 tokens
 else:
 # Capture ZeroDivisionError
 try:
 result = ops_tiles[token](operand_1, operand_2)
 except ZeroDivisionError:
 break
 else:
 # Intermediate numbers can only be whole numbers
 if (result < 0) or (result != int(result)):
 break
 elif result == target:
 postfix = postfix[0:idx + 1]
 break
 else:
 stack.append(result)
 # Else token is a number
 else:
 # Add number to the stack and await an operator
 stack.append(token)
 if result == target:
 break
 else:
 continue
 break
 else:
 continue
 break
 # Print results!
 if result != target:
 print("No solution found!")
 else:
 print(postfix)
 print(postfix_to_infix(postfix))
 # Play again?
 while True:
 try:
 print("Do you want to play again (y/n)?")
 play_again = input()[0]
 except IndexError:
 pass
 else:
 if play_again.lower() == 'n':
 play_new_game = False
 break
 elif play_again.lower() == 'y':
 print("\n==========")
 break
if __name__ == '__main__':
 main()

A short explanation

The above code does the following things:

1. As per the game rules, the player is asked how many 'big' numbers they want (between 0 to 4), and then the 6 numbers are picked out at random (the remaining being 'small' numbers) and printed.
2. The target number (random number between 100 to 999) is printed out as well.
3. It tries to find a solution via brute force (iterating through all possible combinations of numbers and operators to achieve the target number). I use Reverse Polish notation in my solution-searching algorithm.
4. Print the solution, if available, and do so in (standard) Polish notation. The conversion from RPN to PN is via the function postfix_to_infix function.

Notes and Questions

Just for context: I'm relatively new to Python and OOP in general; I have a better grasp of functional programming. I would gladly appreciate some feedback regarding:

• How to make the code more efficient? Currently it takes a very, very long time to report a given combination has no solution. NB: It will be even worse if I allow for more operators in the postfix expression (>4). As such, I am not 100% sure that this code will find the solutions for all possible games, but I can only increase the number of operators once my code runs more efficiently -- otherwise it would lengthen the time to even find solutions of extremely easy games (e.g. 100, 25, 2, 4, 6, 8 with target 650).
• There are indeed some efficiency checks suggested in the link above (by DataGenetics) that I did not implement in my code as I'm not too sure yet (e.g. steps that multiply and divide by 1 -- wasting moves)
• Would it be possible to implement a class for this? I'm not sure what the objects will be in this case. I'll also be willing to hear why it isn't worth implementing this in a class, if you feel this way (I suspect this is the case here).
• Any improvements for the postfix_to_infix function and the way it is implemented in Python?

Thank you!

Answer 1

Like the other answers say, split up your program.

A few nice ways you can split it up, and make it more efficient at the same time is to change the way the sequences are generated.

import operator
from itertools import permutations, combinations_with_replacement, product
DEFAULT_OPERATORS = {
 "+": operator.add,
 "-": operator.sub,
 "*": operator.mul,
 "/": operator.truediv,
}

Check whether the sequence is valid

A sequence is lexically if there are enough numbers to use as operands for the operations. This means that at for any place in the sequence, there should be 1 more number than operator. If you present a number as True and an operator as False, this check can be done like this:

def is_valid(sequence):
 count_nums = 0
 for token in sequence:
 if not token:
 count_nums -= 1
 if count_nums < 1:
 return False
 else:
 count_nums +=1
 return True

assert is_valid((True, True, False))
assert is_valid((True, True, False, True, True, False, False))
assert is_valid((True, True, False, True, False))
assert not is_valid((True, True, False, False))
assert not is_valid((True, False, True, False))
assert not is_valid((True, False))
assert not is_valid((False, True, True, False))

Generate the sequences

If we know how many numbers and operators are used, we can generate all valid sequences like this:

def valid_sequences(num_count):
 base = (True,) * num_count + (False,) * (num_count - 1)
# print(base)
 perms = distinct_permutations(base)
# perms = list(perms)
# print(f'possible permutations: {len(perms)}')
 for sequence in perms:
# print(sequence)
 if is_valid(sequence):
 yield sequence

print(list(valid_sequences(2)))
print(list(valid_sequences(3)))
print(list(valid_sequences(4)))

[(True, True, False)]
[(True, True, True, False, False), (True, True, False, True, False)]
[(True, True, True, True, False, False, False), (True, True, True, False, True, False, False), (True, True, True, False, False, True, False), (True, True, False, True, True, False, False), (True, True, False, True, False, True, False)]

if you look at the case of 6 numbers and 5 operators, there are 462 permutations, but only 42 valid ones.

Populate the sequence

def distinct_ordered_combinations(iterable, r):
 seen = set()
 for combination in combinations_with_replacement(iterable, r):
 for permutation in permutations(combination):
 if permutation not in seen:
 yield permutation
 seen.add(permutation)
def generate_sequences(numbers, operators):
 l = len(numbers)
 for sequence in valid_sequences(l):
 for nums, ops in product(
 distinct_permutations(numbers),
 distinct_ordered_combinations(operators, l - 1),
 ):
 nums = iter(nums)
 ops = iter(ops)
 yield tuple(next(nums) if token else next(ops) for token in sequence)

list(generate_sequences((1, 2), operators=DEFAULT_OPERATORS))

[(1, 2, '+'),
 (1, 2, '-'),
 (1, 2, '*'),
 (1, 2, '/'),
 (2, 1, '+'),
 (2, 1, '-'),
 (2, 1, '*'),
 (2, 1, '/')]

 list(generate_sequences((1, 2, 2), operators='+-'))

[(1, 2, 2, '+', '+'),
 (1, 2, 2, '+', '-'),
 (1, 2, 2, '-', '+'),
 (1, 2, 2, '-', '-'),
 (2, 1, 2, '+', '+'),
 (2, 1, 2, '+', '-'),
 (2, 1, 2, '-', '+'),
 (2, 1, 2, '-', '-'),
 (2, 2, 1, '+', '+'),
 (2, 2, 1, '+', '-'),
 (2, 2, 1, '-', '+'),
 (2, 2, 1, '-', '-'),
 (1, 2, '+', 2, '+'),
 (1, 2, '+', 2, '-'),
 (1, 2, '-', 2, '+'),
 (1, 2, '-', 2, '-'),
 (2, 1, '+', 2, '+'),
 (2, 1, '+', 2, '-'),
 (2, 1, '-', 2, '+'),
 (2, 1, '-', 2, '-'),
 (2, 2, '+', 1, '+'),
 (2, 2, '+', 1, '-'),
 (2, 2, '-', 1, '+'),
 (2, 2, '-', 1, '-')]

The nice thing, is that each of these steps can be verified and tested, which can't be done using your giant nested for-loop

Calculate the intermediate results for the sequences

Since you don't expect trouble with the order of operators anymore, the checking for the IndexError can be dropped. In this case, I think it would be simpler to communicate the non-in or negative intermediate result via a custom Exception.

class FalseOperation(ValueError):
 pass
def calculate_sequence(sequence, operators):
 # Implement postfix algorithm
 stack = list(sequence[:2])
 for idx, token in enumerate(sequence[2:], start=3):
 if token not in operators:
 stack.append(token)
 continue
 operand2 = stack.pop()
 operand1 = stack.pop()
 try:
 result = operators[token](operand1, operand2)
 except ZeroDivisionError:
 raise FalseOperation('Division by zero')
 if int(result) != result:
 raise FalseOperation('non-int value')
 if result < 0:
 raise FalseOperation('negative value')
 yield result, sequence[:idx]
 stack.append(result)

list(calculate_sequence((7, 3, 2, 8, 50, '+', '-', 10, '*', '+', '-'), > DEFAULT_OPERATORS))

[(58, (7, 3, 2, 8, 50, '+')),
 (56, (7, 3, 2, 8, 50, '+', '-')),
 (560, (7, 3, 2, 8, 50, '+', '-', 10, '*')),
 (563, (7, 3, 2, 8, 50, '+', '-', 10, '*', '+')),
 (556, (7, 3, 2, 8, 50, '+', '-', 10, '*', '+', '-'))]

Test the sequences

Since the other methods do the heavy lifting, this becomes rather trivially. All you have to do is supply the target, the sequence of numbers, and how many operators you want mixed in

def test_sequences(target, numbers, operators):
 for sequence in generate_sequences(numbers, operators):
 try:
 for result , seq in calculate_sequence(sequence, operators):
 if result == target:
 yield seq
 except FalseOperation:
 pass

list(test_sequences(56, (2, 8, 50), DEFAULT_OPERATORS))

[(2, 8, 50, '+', '-'),
 (2, 50, 8, '+', '-'),
 (8, 2, 50, '-', '+'),
 (50, 2, 8, '-', '+'),
 (2, 8, '-', 50, '+'),
 (2, 50, '-', 8, '+')]

• Troy: Hi, intriguing solution and review (thanks!). I'm still in the midst of looking through this: but (1) is the count=0 statement in valid_sequences is a remnant from your testing?, and (2) generate_sequences doesn't take an op_count from your original definition, right?
• Maarten Fabré: those were indeed left-overs from previous version. In fixing this, I noticed a bug in the generation code. the op_count is unnecessary, since the number of operators is the number of numbers - 1
• Troy: Sorry if it's really basic, but I can't seem to work out how to implement these: the test_sequences function is not being run (I placed print statements inside the function -- they aren't being printed..) I left my code on pastebin instead of an edit to the question as I suppose that's frowned upon here. Could you go into a bit more detail how to use this? Btw, in the generate_sequences function, the valid_sequences now only takes l as an input.
• Troy: After a short while of testing, I have to assign the function to something for it to be run? a, b = test_sequences(...)? And presumably to use operators[token](operand1, operand2), operators needs to be a dictionary, but operators in the other functions are all tuples?
• Maarten Fabré: the other functions are oblivious whether operators is a tuple, dict or other sequence. If passed a dict, they use dict.keys(). only calculate_sequence needs a dict

Answer 2

Let's start with the postfix_to_infix function. Since you are already using f-strings elsewhere in your code, you could also use them here. Also note that ''.join(block), is the same as block, because it is a single string.

Instead of [...][0], you can simply use next(...), which will only generate the first value and not the whole list only to take the first element. The only other change you need is to catch StopIteration, instead of IndexError, in case there is no operator left.

I would use a set for the operators, so you get \$\mathcal{O}(1)\$ lookup and define it only once instead of on every loop iteration.

soln is not a very nice name. Either spell it out as solution (it is not that long) or use a name that more clearly communicates that this is just the variable you store the output it (maybe out?). The same is true for tkn. I don't even have a clue what it stands for. But you can re-use the name it is assigned to here, op. I would also use the generic i, instead of op_idx, because it is the only index in this method and allows all lines to fit into the 80 character limit recommended by PEP8, Python's official style-guide.

def postfix_to_infix(postfix):
 """
 :param postfix: tuple containing solution in reverse polish notation
 :return: string of solution in polish notation
 """
 operators = set('+-*/')
 out = list(postfix)
 while True:
 try:
 i, op = next((i, op) for i, op in enumerate(out) if op in operators)
 except StopIteration:
 break
 else:
 out[i - 2:i + 1] = [f"({out[i - 2]}{op}{out[i - 1]})"]
 return out

---

Now for your main function. The general structure of a Python module should be the following:

import somemodule
from someothermodule import something
GLOBAL_CONSTANT = None
class ClassName:
 pass
def function():
 pass
if __name__ == "__main__":
 # Some code using all the above or a direct call to `main()`, if it exists
 c = ClassName()
 function()

According to this, you should pull your imports out of the main function.

You should also start splitting your code into smaller functions, so you can re-use them elsewhere if needed. A prime candidate for this is a function that asks the user for input, validating it at the same time:

def ask_user(message, type_=str, allowed=None):
 while True:
 try:
 usr_input = type_(input(message))
 except ValueError:
 print(f"Please enter a(n) '{type_.__name__}'")
 else:
 if allowed is not None and usr_input not in allowed:
 print(f"Value must be in {allowed}")
 continue
 return usr_input

This function can be used for both of your user inputs:

usr_large_num = ask_user("How many LARGE numbers do you want (0-4)? ", int, range(5))
play_new_game = ask_user("Do you want to play again (y/n)?", allowed="yYnN").lower() == "y"

This has slightly less functionality than your code ("yes" is no longer recognized as "y", a small price to pay, though, IMO).

• Troy: Thanks for the feedback, I learnt a few things from your answer (for one, it did not occur to me that f-strings could be used outside the print function -- it seems obvious now..). To clear things up, tkn stands for token, but I agree wholeheartedly with your comments. Cheers.

Answer 3

The main function is very long: it's probably worth at least pulling out the solution search routine. But beyond that I'm not going to comment much on style, because I'm not a Pythonista and others can do it better than me.

How to make the code more efficient?

Simple: don't repeat yourself. Here I'm not referring to copy-paste code, but to repeated calculations. How many permutations beginning num num op op is the code generating and then discarding with an IndexError on the second op? I make it almost five million.

If the numbers include 2 and 3, how many permutations beginning 2 3 + is the code generating? How many beginning 3 2 + is it generating? How many does it need to generate?

So two keys to making it more efficient are to build the stacks up incrementally such that you only push an operator when there are at least two numbers on the stack; and to use a dictionary or some similar mechanism to identify expressions with the same numbers and the same resulting stack, so that only one of them is combined as part of larger expressions. One approach would be to only work with closed expressions (i.e. those which give a stack of one number) and to only combine them in descending order (since that's what we want for division and subtraction).

As a bonus, this fixes something which IMO is a bug in the original program: Countdown doesn't require all 6 numbers to be used.