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, '+')]

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).

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.