tree.py 10.9 KB
Newer Older
mrk022's avatar
mrk022 committed
1 2
# base tree code from http://knuth.luther.edu/~leekent/CS2Plus/chap6/chap6.html
# changed this code to allow for variable functions, and change infix expressions to postfix
3
# added variable node class to consider when both nodes have variables in them
mrk022's avatar
mrk022 committed
4 5 6

import queue

7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
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 evaluate(self, parent, left):
31
        # can simplify this to call helper function
32 33
        if left:
            if parent.left.canEval():
34 35
                print("parent.left.left is ", parent.left.left.value)
                print("parent.left.right is ", parent.left.right)
36
                if isinstance(parent.left.left, NumNode):
37
                    print("num node")
38 39 40
                    print("evaluating", parent.left.eval())
                    newNode = NumNode(parent.left.eval())
                else:
41
                    print("variable node")
42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87
                    print("evaluating", parent.left.eval())
                    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):
                    print("evaluating", parent.right.eval())
                    newNode = NumNode(parent.right.eval())
                else:
                    print("evaluating", parent.right.eval())
                    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's avatar
mrk022 committed
88 89 90 91
class TimesNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
92
        self.value = "*"
mrk022's avatar
mrk022 committed
93
        
94
    
mrk022's avatar
mrk022 committed
95
    def eval(self):  
96 97 98 99 100 101 102
        '''if check_if_variables(str(self.left.eval()), str(self.right.eval())) == "both_vars":
            v1 = str(self.left.eval()).split("x")[0]
            v2 = str(self.right.eval()).split("x")[0]
            evalStr = v1 + "*" + v2
            return eval(evalStr)
        else:
            return self.left.eval() * self.right.eval()'''
mrk022's avatar
mrk022 committed
103 104 105 106 107
        #if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
        #    return self.left.eval()*self.right.eval()
        #elif isinstance(self.left,VariableNode) or isinstance(self.right,VariableNode):
        #    print("can't combine variables of differing types")
            #need some sort of combined node here i think
108
        #else:
mrk022's avatar
mrk022 committed
109 110 111 112 113 114 115 116 117
        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
118
        self.value = "+"
mrk022's avatar
mrk022 committed
119
        
120
    def canEval(self):
121
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
122
            return True
123 124 125 126
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
            return True
        else:
            return False
127 128 129 130

    def getInverse(self):
        return MinusNode(None,None)

mrk022's avatar
mrk022 committed
131
    def eval(self):
132 133 134
        #print("left value ", self.left)
        #print(isinstance(self.left, NumNode))
        #print("right value ", self.right.value())
mrk022's avatar
mrk022 committed
135 136 137 138 139 140 141 142 143 144
        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
145 146 147
        self.value = "-"

    def canEval(self):
148 149 150
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
            return True
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
151
            return True
152 153
        else:
            return False
154 155 156

    def getInverse(self):
        return PlusNode(None, None)
mrk022's avatar
mrk022 committed
157 158 159 160 161 162 163 164 165

    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):
166 167 168
        self.value = num
        self.left = None
        self.right = None
mrk022's avatar
mrk022 committed
169 170
        
    def eval(self):
171 172 173 174
        return self.value

    def canEval(self):
        return False # not an operator
mrk022's avatar
mrk022 committed
175 176
    
    def inorder(self):
177
        return str(self.value)
mrk022's avatar
mrk022 committed
178

179 180 181
class VariableNode:
    def __init__(self, value):
        self.value = value
182 183
        self.left = None
        self.right = None
184 185 186

    def eval(self):
        # get value before x
187 188 189 190
        return int(self.value.split("x")[0])

    def canEval(self):
        return False #not an operator
191 192 193 194

    def inorder(self):
        return str(self.eval())

mrk022's avatar
mrk022 committed
195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216
def check_if_variables(x, y):
    if 'x' in x and 'x' in y:
        return "both_vars"
    elif 'x' in x or 'x' in y:
        return "one_var"
    else:
        return "both_const"

def E(q):
    if q.isEmpty():
        raise ValueError("Invalid Prefix Expression")
    
    token = q.dequeue()
    if token == "+":
        return PlusNode(E(q),E(q))
    
    if token == "*":
        return TimesNode(E(q),E(q))

    if token == "-":
        return MinusNode(E(q), E(q))
    
217 218 219 220 221
    #try:
    if "x" in token:
        output = VariableNode(token)
        # make a different class for variable node
    else:
222
        output = NumNode(int(token))
223 224 225
    #except:
    #    print("Not a valid operation, need variables of like terms")
    #    output = NumNode(float(0))
mrk022's avatar
mrk022 committed
226 227 228 229 230 231
    return output

OPERATORS = set(['+', '-', '*', '/', '(', ')'])
PRIORITY = {'+':1, '-':1, '*':2, '/':2}

def infix_to_prefix(formula):
mrk022's avatar
mrk022 committed
232
    # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py
mrk022's avatar
mrk022 committed
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
    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)
    
    # leftover
    while op_stack:
        op = op_stack.pop()
        a = exp_stack.pop()
        b = exp_stack.pop()
        exp_stack.append( op+" "+b+" "+a )
261
    print(exp_stack[-1])
mrk022's avatar
mrk022 committed
262 263
    return exp_stack[-1]

264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316
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
    
mrk022's avatar
mrk022 committed
317 318
  
def eval_tree(expression):
mrk022's avatar
mrk022 committed
319
    '''x = NumNode(5)
mrk022's avatar
mrk022 committed
320 321 322 323 324 325
    y = NumNode(4)
    p = PlusNode(x,y)
    t = TimesNode(p, NumNode(6))
    root = PlusNode(t, NumNode(3))
    
    print(root.eval())
mrk022's avatar
mrk022 committed
326
    print(root.inorder())'''
mrk022's avatar
mrk022 committed
327 328 329 330 331 332 333 334 335 336 337 338
    
    #x = input("Please enter a prefix expression: ")
    x = infix_to_prefix(expression)

    lst = x.split()
    q = queue.Queue()
    
    for token in lst:
        q.enqueue(token)
        
    root = E(q)
    
mrk022's avatar
mrk022 committed
339
    result = root.eval()
mrk022's avatar
mrk022 committed
340 341
    print(root.eval())
    print(root.inorder())
mrk022's avatar
mrk022 committed
342
    return result
mrk022's avatar
mrk022 committed
343 344 345 346 347 348 349
    #print(infix_to_prefix("3 + 4 - 6"))

def split_equation(equation):
    expressions = equation.split("=")
    return expressions

def main():
350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368
    #exp = split_equation("4x * 3x + 2")
    #x = eval_tree(exp[0])
    '''eqTree = EquationTree()
    x1 = VariableNode("4x")
    x2 = NumNode(3)
    x3 = NumNode(2)
    x4 = NumNode(7)
    t = PlusNode(x1,x2)
    t2 = PlusNode(x3,x4)
    eqTree.root.left = t
    eqTree.root.right = t2
    eqTree.printTree()'''
    equationTree = generate_tree("4x + 3 = 2 + 7")
    print(equationTree.root.left.canEval())
    print(equationTree.root.right.canEval())
    newNode = NumNode(equationTree.root.right.eval())
    equationTree.updateNode(equationTree.root, False,  newNode)
    equationTree.printTree()
    equationTree.get_expression()
mrk022's avatar
mrk022 committed
369
        
370 371 372 373 374 375

'''
print("old node should be 2, is ", oldNum.value)
                print("moved node should be 7, is ", moveNode.value)
                print("new op is ", newOp.value)
                '''
mrk022's avatar
mrk022 committed
376 377 378
    
if __name__ == "__main__":
    main()