-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathtree.py
54 lines (48 loc) · 1.47 KB
/
tree.py
1
2
3
4
5
6
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
import argparse
class TreeNode:
def __init__(self, val, left=None, right=None):
self.val = val
self.left = left
self.right = right
def __repr__(self):
return 'TreeNode({})'.format(self.val)
def deserialize(string):
if string == '{}':
return None
nodes = [None if val == 'null' else TreeNode(int(val))
for val in string.strip('[]{}').split(',')]
kids = nodes[::-1]
root = kids.pop()
for node in nodes:
if node:
if kids: node.left = kids.pop()
if kids: node.right = kids.pop()
return root
def drawtree(root):
def height(root):
return 1 + max(height(root.left), height(root.right)) if root else -1
def jumpto(x, y):
t.penup()
t.goto(x, y)
t.pendown()
def draw(node, x, y, dx):
if node:
t.goto(x, y)
jumpto(x, y-20)
t.write(node.val, align='center', font=('Arial', 12, 'normal'))
draw(node.left, x-dx, y-60, dx/2)
jumpto(x, y-20)
draw(node.right, x+dx, y-60, dx/2)
import turtle
t = turtle.Turtle()
t.speed(0); turtle.delay(0)
h = height(root)
jumpto(0, 30*h)
draw(root, 0, 30*h, 40*h)
t.hideturtle()
turtle.mainloop()
parser = argparse.ArgumentParser(description="xxx")
parser.add_argument('tree', help='The tree representation')
if __name__ == '__main__':
args = parser.parse_args()
drawtree(deserialize(args.tree))