tree.py 8.53 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
23
24
25
26
27
28
29
30
31
32
33
34
35
36
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)

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
        if left:
            if parent.left.canEval():
54
55
                print("parent.left.left is ", parent.left.left.value)
                print("parent.left.right is ", parent.left.right)
56
                if isinstance(parent.left.left, NumNode):
57
                    print("num node")
58
59
60
                    print("evaluating", parent.left.eval())
                    newNode = NumNode(parent.left.eval())
                else:
61
                    print("variable node")
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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
                    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
107
108
109
110
class TimesNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
111
        self.value = "*"
mrk022's avatar
mrk022 committed
112
        
113
    
mrk022's avatar
mrk022 committed
114
115
116
117
118
119
120
121
122
123
    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
124
        self.value = "+"
mrk022's avatar
mrk022 committed
125
        
126
    def canEval(self):
127
        if isinstance(self.left, VariableNode) and isinstance(self.right,VariableNode):
128
            return True
129
130
131
132
        elif isinstance(self.left, NumNode) and isinstance(self.right, NumNode):
            return True
        else:
            return False
133
134
135
136

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

mrk022's avatar
mrk022 committed
137
138
139
140
141
142
143
144
145
146
147
    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
148
149
150
        self.value = "-"

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

    def getInverse(self):
        return PlusNode(None, None)
mrk022's avatar
mrk022 committed
160
161
162
163
164
165
166
167
168

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

    def canEval(self):
        return False # not an operator
mrk022's avatar
mrk022 committed
178
179
    
    def inorder(self):
180
        return str(self.value)
mrk022's avatar
mrk022 committed
181

182
183
184
class VariableNode:
    def __init__(self, value):
        self.value = value
185
186
        self.left = None
        self.right = None
187
188
189

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

    def canEval(self):
mrk022's avatar
mrk022 committed
193
        return False 
194
195
196
197

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

mrk022's avatar
mrk022 committed
198
199
200
201
OPERATORS = set(['+', '-', '*', '/', '(', ')'])
PRIORITY = {'+':1, '-':1, '*':2, '/':2}

def infix_to_prefix(formula):
mrk022's avatar
mrk022 committed
202
    # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py
mrk022's avatar
mrk022 committed
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
    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 )
229
    print(exp_stack[-1])
mrk022's avatar
mrk022 committed
230
231
    return exp_stack[-1]

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
274
275
276
277
278
279
280
281
282
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()
mrk022's avatar
mrk022 committed
283
    return t