Tries to solve the problem of: 1 1 1 = 6, 2 2 2 = 6, ...

pub_puzzle.py

#!/usr/bin/env python3
#
# Usage: pub.py [NUM]
# Tries to solve the problem of:
# 1 1 1 = 6
# 2 2 2 = 6
# 3 3 3 = 6
# ...

import sys
import re
import random
from random import choice

import ast
import operator as op
import math
import functools
import numbers


sqrt = lambda a: math.sqrt(a)
fact = lambda a: math.factorial(a)
neg = lambda a: -a
div = lambda a, b: a/b
powr = lambda a, b: pow(a, b)

DOUBLE_OPS = op.add, op.sub, op.mul, div, powr #, op.floordiv , op.mod
SINGLE_OPS = op.abs, sqrt, fact, neg

OP_STR = {sqrt: 'sqrt', fact: '!', neg: '-', div: '/', op.add: '+', op.sub: '-',
    op.mul: '*', op.abs: 'abs', powr: 'pwr'}

class Node:
    def __init__(self, is_root, left, right):
        self.is_root = is_root
        self.single_ops = []
        self.op = [op.add]
        self.left = left
        self.right = right
    def get_val(self):
        try:
            left_val = self.left.get_val()
            right_val = self.right.get_val()
            if not left_val or not right_val:
                return
            try:
                combined = self.op[0](left_val, right_val)
            except OverflowError:
                return
            if not combined:
                return
            return run_pipe(combined, self.single_ops)

        except ValueError:
            return
    def __repr__(self):
        out = f'{ops_to_str(self.single_ops)}({self.left} {OP_STR[self.op[0]]} {self.right})'
        if self.is_root:
            out += f' = {self.get_val()}'
        return out


def run_pipe(val, functions):
    for fun in functions:
        if val is None or abs(val) > 1000:
            return None
        val = fun(val)
    return val


def ops_to_str(ops):
    if not ops:
        return ''
    if len(ops) == 1:
        return f'{OP_STR[ops[0]]} '
    return f'{[OP_STR[a] for a in ops]} '


class Leaf:
    def __init__(self, val):
        self.val = val
        self.single_ops = []
    def get_val(self):
        return run_pipe(self.val, self.single_ops)
    def __repr__(self):
        return f'{ops_to_str(self.single_ops)}{self.val}'


def main():
    if len(sys.argv) < 2:
        for i in range(1, 10):
            print(run(i))
    else:
        print(run(int(sys.argv[1])))


def run(a_sum):
    tree, params = get_tree(a_sum)
    while tree.get_val() != 6:
        val = tree.get_val()
        if val is None or abs(val) > 1000 or abs(val) < 0.01:
            tree, params = get_tree(a_sum)
            continue
        target_param = random.choice(params)
        if len(target_param) == 1 and target_param[0] in DOUBLE_OPS:
            target_param[0] = random.choice(DOUBLE_OPS)
            continue
        if not target_param:
            target_param.append(random.choice(SINGLE_OPS))
            continue
        dice = random.randint(0, 2)
        if dice == 0:
            target_param.append(choice(SINGLE_OPS))
        elif dice == 1:
            target_param.pop()
        else:
            target_param[random.randint(0, len(target_param)-1)] = choice(SINGLE_OPS)
    return tree


def get_tree(val):
    leaf_1, leaf_2, leaf_3 = Leaf(val), Leaf(val), Leaf(val)
    l_node = Node(False, leaf_1, leaf_2)
    if random.randint(0, 1) == 0:
        root = Node(True, l_node, leaf_3)
    else:    
        root = Node(True, leaf_3, l_node)
    params = root.single_ops, root.op, leaf_1.single_ops, leaf_1.single_ops, \
        leaf_1.single_ops, l_node.single_ops, l_node.op
    return root, params


if __name__ == '__main__':
    main()