''' AICogSci Final Project base tree code from http://knuth.luther.edu/~leekent/CS2Plus/chap6/chap6.html Added: - EquationTree class to generate a specialized tree - Branch and Equals Node classes to make printing more legible - canEval, getInverse and inorder methods to most of the nodes to allow further functionality - VariableNode class was added to account for nodes with variables - add_to_tree and generate_tree methods used to create and change equation tree class''' import queue class EquationTree: def __init__(self): self.root = EqualsNode() self.root.left = None self.root.right = None def updateNode(self, parent, left, newNode): if left: parent.left = newNode else: parent.right = newNode #self.printTree() def get_expression(self): self.output = "" self.inorder(self.root) #print("output is: ", self.output) return self.output def inorder(self, node): if node: self.inorder(node.left) self.output += str(node.value) self.inorder(node.right) def solve(self): # solves equation assuming there are only a variable node and num node left if isinstance(self.root.left, VariableNode): # variablenode is on left of equation sol = self.root.right.eval() / self.root.left.eval() self.root.left = VariableNode('x') self.root.right = NumNode(int(sol)) else: sol = self.root.left.eval() / self.root.right.eval() self.root.right = VariableNode('x') self.root.left = NumNode(int(sol)) return self.get_expression() def evaluate(self, parent, left): # can simplify this to call helper function if left: if parent.left.canEval(): if isinstance(parent.left.left, NumNode): newNode = NumNode(parent.left.eval()) else: newNode = VariableNode(str(parent.left.eval()) + "x") parent.left = newNode elif parent.left.left != None and parent.left.left.canEval(): self.evaluate(parent.left, True) elif parent.left.right != None and parent.left.right.canEval(): self.evaluate(parent.left, False) else: if parent.right.canEval(): if isinstance(parent.right.left, NumNode): newNode = NumNode(parent.right.eval()) else: newNode = VariableNode(str(parent.right.eval()) + "x") parent.right = newNode elif parent.right.left != None and parent.right.left.canEval(): self.evaluate(parent.right, True) elif parent.right.right != None and parent.right.right.canEval(): self.evaluate(parent.right, False) def printTree(self): q = queue.Queue() q.enqueue(self.root) while not q.isEmpty(): currentNode = q.front() if currentNode.left != None: b = Branch() q.enqueue(b) q.enqueue(currentNode.left) if currentNode.right != None: q.enqueue(currentNode.right) print(q.dequeue().value, end =' ') class Branch: def __init__(self): self.value = "\n" self.left = None self.right = None class EqualsNode: def __init__(self): self.value = "=" class TimesNode: def __init__(self, left, right): self.left = left self.right = right self.value = "*" def eval(self): return self.left.eval() * self.right.eval() def inorder(self): return "(" + self.left.inorder() + " * " + self.right.inorder() + ")" class PlusNode: def __init__(self, left, right): self.left = left self.right = right self.value = "+" def canEval(self): if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode): return True elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode): return True else: return False def getInverse(self): return MinusNode(None,None) def eval(self): return self.left.eval() + self.right.eval() def inorder(self): return "(" + self.left.inorder() + " + " + self.right.inorder() + ")" class MinusNode: def __init__(self,left, right): self.left = left self.right = right self.value = "-" def canEval(self): if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode): return True elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode): return True else: return False def getInverse(self): return PlusNode(None, None) def eval(self): return self.left.eval() - self.right.eval() def inorder(self): return "(" + self.left.inorder() + "-" + self.right.inorder() + ")" class NumNode: def __init__(self, num): self.value = num self.left = None self.right = None def eval(self): return self.value def canEval(self): return False # not an operator def inorder(self): return str(self.value) class VariableNode: def __init__(self, value): self.value = value self.left = None self.right = None def eval(self): # get value before x return int(self.value.split("x")[0]) def canEval(self): return False def inorder(self): return str(self.eval()) OPERATORS = set(['+', '-', '*', '/', '(', ')']) PRIORITY = {'+':1, '-':1, '*':2, '/':2} def infix_to_prefix(formula): # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py op_stack = [] exp_stack = [] for ch in formula.split(): if not ch in OPERATORS: exp_stack.append(ch) elif ch == '(': op_stack.append(ch) elif ch == ')': while op_stack[-1] != '(': op = op_stack.pop() a = exp_stack.pop() b = exp_stack.pop() exp_stack.append( op+b+a ) op_stack.pop() # pop '(' else: while op_stack and op_stack[-1] != '(' and PRIORITY[ch] <= PRIORITY[op_stack[-1]]: op = op_stack.pop() a = exp_stack.pop() b = exp_stack.pop() exp_stack.append( op+" "+b+" "+a ) op_stack.append(ch) while op_stack: op = op_stack.pop() a = exp_stack.pop() b = exp_stack.pop() exp_stack.append( op+" "+b+" "+a ) #print(exp_stack[-1]) return exp_stack[-1] def add_to_tree(q, root, left): if not q.isEmpty(): token = q.dequeue() operator = False if token == "+": newNode = PlusNode(None, None) operator = True elif token == "*": newNode = TimesNode(None, None) operator = True elif token == "-": newNode = MinusNode(None, None) operator = True elif "x" in token: newNode = VariableNode(token) else: newNode = NumNode(int(token)) if left: root.left = newNode else: root.right = newNode if operator: add_to_tree(q, newNode, True) add_to_tree(q, newNode, False) return newNode def generate_tree(equation): t = EquationTree() expressions = equation.split("=") x1 = infix_to_prefix(expressions[0]) x2 = infix_to_prefix(expressions[1]) lst1 = x1.split() q1 = queue.Queue() lst2 = x2.split() q2 = queue.Queue() for token in lst1: q1.enqueue(token) add_to_tree(q1, t.root, True) # add expression tree to left side of equation tree for token in lst2: q2.enqueue(token) add_to_tree(q2, t.root, False) # add expression tree to right side of equation tree #t.printTree() return t