tree.py 10.1 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
31
32
33
34
35
36
37
38
39
40
41
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
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):
        if left:
            if parent.left.canEval():
                if isinstance(parent.left.left, NumNode):
                    print("evaluating", parent.left.eval())
                    newNode = NumNode(parent.left.eval())
                else:
                    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
83
84
85
86
class TimesNode:
    def __init__(self, left, right):
        self.left = left
        self.right = right
87
        self.value = "*"
mrk022's avatar
mrk022 committed
88
        
89
    
mrk022's avatar
mrk022 committed
90
    def eval(self):  
91
92
93
94
95
96
97
        '''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
98
99
100
101
102
        #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
103
        #else:
mrk022's avatar
mrk022 committed
104
105
106
107
108
109
110
111
112
        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
113
        self.value = "+"
mrk022's avatar
mrk022 committed
114
        
115
116
117
118
119
120
121
122
    def canEval(self):
        if isinstance(self.left, VariableNode) == isinstance(self.right,VariableNode):
            return True
        return False

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

mrk022's avatar
mrk022 committed
123
124
125
126
127
128
129
130
131
132
133
    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
134
135
136
137
138
139
140
141
142
        self.value = "-"

    def canEval(self):
        if isinstance(self.left, VariableNode) == isinstance(self.right,VariableNode):
            return True
        return False

    def getInverse(self):
        return PlusNode(None, None)
mrk022's avatar
mrk022 committed
143
144
145
146
147
148
149
150
151

    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):
152
153
154
        self.value = num
        self.left = None
        self.right = None
mrk022's avatar
mrk022 committed
155
156
        
    def eval(self):
157
158
159
160
        return self.value

    def canEval(self):
        return False # not an operator
mrk022's avatar
mrk022 committed
161
162
    
    def inorder(self):
163
        return str(self.value)
mrk022's avatar
mrk022 committed
164

165
166
167
class VariableNode:
    def __init__(self, value):
        self.value = value
168
169
        self.left = None
        self.right = None
170
171
172

    def eval(self):
        # get value before x
173
174
175
176
        return int(self.value.split("x")[0])

    def canEval(self):
        return False #not an operator
177
178
179
180

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

mrk022's avatar
mrk022 committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
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))
    
203
204
205
206
207
    #try:
    if "x" in token:
        output = VariableNode(token)
        # make a different class for variable node
    else:
208
        output = NumNode(int(token))
209
210
211
    #except:
    #    print("Not a valid operation, need variables of like terms")
    #    output = NumNode(float(0))
mrk022's avatar
mrk022 committed
212
213
214
215
216
217
    return output

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

def infix_to_prefix(formula):
mrk022's avatar
mrk022 committed
218
    # this method was taken from open source https://github.com/lilianweng/LeetcodePython/blob/master/expression.py
mrk022's avatar
mrk022 committed
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    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 )
247
    print(exp_stack[-1])
mrk022's avatar
mrk022 committed
248
249
    return exp_stack[-1]

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
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
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
303
304
  
def eval_tree(expression):
mrk022's avatar
mrk022 committed
305
    '''x = NumNode(5)
mrk022's avatar
mrk022 committed
306
307
308
309
310
311
    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
312
    print(root.inorder())'''
mrk022's avatar
mrk022 committed
313
314
315
316
317
318
319
320
321
322
323
324
    
    #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
325
    result = root.eval()
mrk022's avatar
mrk022 committed
326
327
    print(root.eval())
    print(root.inorder())
mrk022's avatar
mrk022 committed
328
    return result
mrk022's avatar
mrk022 committed
329
330
331
332
333
334
335
    #print(infix_to_prefix("3 + 4 - 6"))

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

def main():
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
    #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
355
356
357
358
        
    
if __name__ == "__main__":
    main()