tree.py 8.1 KB
 mrk022 committed Dec 07, 2019 1 2 3 4 5 6 7 8 9 10 ``````''' 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''' `````` mrk022 committed Nov 28, 2019 11 12 13 `````` import queue `````` mrk022 committed Dec 04, 2019 14 15 16 17 18 19 20 21 22 ``````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 `````` mrk022 committed Dec 08, 2019 23 `````` #self.printTree() `````` mrk022 committed Dec 04, 2019 24 25 26 27 `````` def get_expression(self): self.output = "" self.inorder(self.root) `````` mrk022 committed Dec 08, 2019 28 `````` #print("output is: ", self.output) `````` mrk022 committed Dec 04, 2019 29 30 31 32 33 34 35 36 `````` return self.output def inorder(self, node): if node: self.inorder(node.left) self.output += str(node.value) self.inorder(node.right) `````` mrk022 committed Dec 07, 2019 37 38 39 40 41 42 43 44 45 46 47 48 49 `````` 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() `````` mrk022 committed Dec 04, 2019 50 `````` def evaluate(self, parent, left): `````` mrk022 committed Dec 06, 2019 51 `````` # can simplify this to call helper function `````` mrk022 committed Dec 04, 2019 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 `````` 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()) `````` mrk022 committed Dec 08, 2019 67 `````` else: `````` mrk022 committed Dec 04, 2019 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 `````` 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 = "=" `````` mrk022 committed Nov 28, 2019 99 100 101 102 ``````class TimesNode: def __init__(self, left, right): self.left = left self.right = right `````` mrk022 committed Dec 04, 2019 103 `````` self.value = "*" `````` mrk022 committed Nov 28, 2019 104 `````` `````` mrk022 committed Dec 04, 2019 105 `````` `````` mrk022 committed Nov 28, 2019 106 107 108 109 110 111 112 113 114 115 `````` 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 `````` mrk022 committed Dec 04, 2019 116 `````` self.value = "+" `````` mrk022 committed Nov 28, 2019 117 `````` `````` mrk022 committed Dec 04, 2019 118 `````` def canEval(self): `````` mrk022 committed Dec 06, 2019 119 `````` if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode): `````` mrk022 committed Dec 04, 2019 120 `````` return True `````` mrk022 committed Dec 06, 2019 121 122 123 124 `````` elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode): return True else: return False `````` mrk022 committed Dec 04, 2019 125 126 127 128 `````` def getInverse(self): return MinusNode(None,None) `````` mrk022 committed Nov 28, 2019 129 130 131 132 133 134 135 136 137 138 139 `````` 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 `````` mrk022 committed Dec 04, 2019 140 141 142 `````` self.value = "-" def canEval(self): `````` mrk022 committed Dec 06, 2019 143 144 145 `````` if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode): return True elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode): `````` mrk022 committed Dec 04, 2019 146 `````` return True `````` mrk022 committed Dec 06, 2019 147 148 `````` else: return False `````` mrk022 committed Dec 04, 2019 149 150 151 `````` def getInverse(self): return PlusNode(None, None) `````` mrk022 committed Nov 28, 2019 152 153 154 155 156 157 158 159 160 `````` 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): `````` mrk022 committed Dec 04, 2019 161 162 163 `````` self.value = num self.left = None self.right = None `````` mrk022 committed Nov 28, 2019 164 165 `````` def eval(self): `````` mrk022 committed Dec 04, 2019 166 167 168 169 `````` return self.value def canEval(self): return False # not an operator `````` mrk022 committed Nov 28, 2019 170 171 `````` def inorder(self): `````` mrk022 committed Dec 04, 2019 172 `````` return str(self.value) `````` mrk022 committed Nov 28, 2019 173 `````` `````` mrk022 committed Dec 02, 2019 174 175 176 ``````class VariableNode: def __init__(self, value): self.value = value `````` mrk022 committed Dec 04, 2019 177 178 `````` self.left = None self.right = None `````` mrk022 committed Dec 02, 2019 179 180 181 `````` def eval(self): # get value before x `````` mrk022 committed Dec 04, 2019 182 183 184 `````` return int(self.value.split("x")[0]) def canEval(self): `````` mrk022 committed Dec 07, 2019 185 `````` return False `````` mrk022 committed Dec 02, 2019 186 187 188 189 `````` def inorder(self): return str(self.eval()) `````` mrk022 committed Nov 28, 2019 190 191 192 193 ``````OPERATORS = set(['+', '-', '*', '/', '(', ')']) PRIORITY = {'+':1, '-':1, '*':2, '/':2} def infix_to_prefix(formula): `````` mrk022 committed Dec 02, 2019 194 `````` # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py `````` mrk022 committed Nov 28, 2019 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 `````` 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 ) `````` mrk022 committed Dec 08, 2019 221 `````` #print(exp_stack[-1]) `````` mrk022 committed Nov 28, 2019 222 223 `````` return exp_stack[-1] `````` mrk022 committed Dec 04, 2019 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 ``````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 `````` mrk022 committed Dec 08, 2019 274 `````` #t.printTree() `````` mrk022 committed Dec 07, 2019 275 `` return t``