tree.py 8.1 KB
Newer Older
mrk022's avatar
mrk022 committed
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's avatar
mrk022 committed
11 12 13

import queue

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's avatar
mrk022 committed
23
        #self.printTree()
24 25 26 27

    def get_expression(self):
        self.output = ""
        self.inorder(self.root)
mrk022's avatar
mrk022 committed
28
        #print("output is: ", self.output)
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's avatar
mrk022 committed
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()

50
    def evaluate(self, parent, left):
51
        # can simplify this to call helper function
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's avatar
mrk022 committed
67
                else:                  
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's avatar
mrk022 committed
99 100 101 102
class TimesNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
103
        self.value = "*"
mrk022's avatar
mrk022 committed
104
        
105
    
mrk022's avatar
mrk022 committed
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
116
        self.value = "+"
mrk022's avatar
mrk022 committed
117
        
118
    def canEval(self):
119
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
120
            return True
121 122 123 124
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
            return True
        else:
            return False
125 126 127 128

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

mrk022's avatar
mrk022 committed
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
140 141 142
        self.value = "-"

    def canEval(self):
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):
146
            return True
147 148
        else:
            return False
149 150 151

    def getInverse(self):
        return PlusNode(None, None)
mrk022's avatar
mrk022 committed
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):
161 162 163
        self.value = num
        self.left = None
        self.right = None
mrk022's avatar
mrk022 committed
164 165
        
    def eval(self):
166 167 168 169
        return self.value

    def canEval(self):
        return False # not an operator
mrk022's avatar
mrk022 committed
170 171
    
    def inorder(self):
172
        return str(self.value)
mrk022's avatar
mrk022 committed
173

174 175 176
class VariableNode:
    def __init__(self, value):
        self.value = value
177 178
        self.left = None
        self.right = None
179 180 181

    def eval(self):
        # get value before x
182 183 184
        return int(self.value.split("x")[0])

    def canEval(self):
mrk022's avatar
mrk022 committed
185
        return False 
186 187 188 189

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

mrk022's avatar
mrk022 committed
190 191 192 193
OPERATORS = set(['+', '-', '*', '/', '(', ')'])
PRIORITY = {'+':1, '-':1, '*':2, '/':2}

def infix_to_prefix(formula):
mrk022's avatar
mrk022 committed
194
    # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py
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 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's avatar
mrk022 committed
221
    #print(exp_stack[-1])
mrk022's avatar
mrk022 committed
222 223
    return exp_stack[-1]

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's avatar
mrk022 committed
274
    #t.printTree()
mrk022's avatar
mrk022 committed
275
    return t