利用pygraphviz绘制二叉树
用法:
python treeWriter.py #输出结果到tree.png
或者
python treeWriter.py <output_file_name>
注意需要安装 graphviz 和 pygraphviz,安装方法见前一篇随笔。
当前的实现需要用户交互输入二叉树,默认输入-1作为子树空。
如对应上面的图,可输入
1 2 7 –1 –1 8 –1 –1 –1 3 4 1 –1 -1 2 –1 9 –1 –1 4 5 3 –1 –1 –1 6 –1 –1
可以有更好的输入方法,待改进。
以前写过利用pyqt显示二叉树的随笔,当时用的已经写好的c++实现的二叉树及相应函数,用swig构造一个python可调用的二叉树模块,这里也可以用同样的方法,只要将treeWriter中相应 如 root.left, root.elem改成相应的c++函数转换好的python可调用的二叉树模块中的接口函数名即可。
而且程序有很好的可扩张性,借助graphviz方便而强大的布局功能,可选的顶点,边参数,用户可以方便的通过继承定义自己的tree writer来绘制不同的二叉树,如下图绘制了一颗建立好的huffman tree.
>>> l
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm']
>>> w
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
来一个复杂一点,我利用绘制的二叉树,来判断压缩过程中生成的huffmantree,以及解压缩过程中写文件后(序列化)之后从文件读出重构的huffmantree是否一致。由于’\n’没有处理好:)暂时不显示key value了,只对比下树形。
1.binaryTree.py
简单定义了二叉链表节点,和二叉树类,给出了通过用户交互输入生成二叉树的方法,及中序遍历方法。
2 #support reference the same as c++ not pointer to pointer in python I think
3
4 #Author goldenlock_pku
5 '''
6 Ilustrate the create binary tree with user input
7 inorder travel
8 '''
9 class Node():
10 '''
11 node will be used in BinaryTree
12 '''
13 def __init__(self,elem, left = None, right = None):
14 self.elem = elem
15 self.left = left
16 self.right = right
17 def __str__(self):
18 return str(self.elem)
19
20 class BinaryTree():
21 def __init__(self,root = None):
22 self.root = root
23
24 def createTree(self, endMark = -1):
25 '''create a binary tree with user input'''
26 def createTreeHelp(endMark = -1):
27 elem = raw_input();
28 if (elem == str(endMark)):
29 return None
30 root = Node(elem)
31 root.left = createTreeHelp(endMark)
32 root.right = createTreeHelp(endMark)
33 return root
34 print('Please input the binary tree data wit ' + str(endMark)\
35 +' as endmark')
36 elem = raw_input()
37 self.root = Node(elem)
38 self.root.left = createTreeHelp(endMark)
39 self.root.right = createTreeHelp(endMark)
40
41 def inorderTravel(self):
42 def inorderTravelHelp(root):
43 if not root:
44 return
45 inorderTravelHelp(root.left)
46 print(root)
47 inorderTravelHelp(root.right)
48 inorderTravelHelp(self.root)
2.treeWriter.py
定义了利用pygraphviz生成dot 文件并进一步生成二叉树图案输出到文件。
绘制过程采用递归,深度优先遍历,利用了不可见节点和边从而能够始终确保能够分清是左子树还是又子树,如果不采用
加入不可见节点,边的方法则无法分清,如果绘制树的话可以采用简单的不加不可见节点的方法。
2 from binaryTree import *
3 import pygraphviz as pgv
4 '''
5 treeWriter with func wite can write a binary tree to tree.png or user spcified
6 file
7 '''
8 class treeWriter():
9 def __init__(self, tree):
10 self.num = 1 #mark each visible node as its key
11 self.num2 = -1 #makk each invisible node as its key
12 self.tree = tree
13
14 def write(self, outfile = 'tree.png'):
15 def writeHelp(root, A):
16 if not root:
17 return
18
19 p = str(self.num)
20 self.num += 1
21 A.add_node(p, label = str(root.elem))
22 q = None
23 r = None
24
25 if root.left:
26 q = writeHelp(root.left, A)
27 A.add_node(q, label = str(root.left.elem))
28 A.add_edge(p,q)
29 if root.right:
30 r = writeHelp(root.right, A)
31 A.add_node(r, label = str(root.right.elem))
32 A.add_edge(p,r)
33
34 if q or r:
35 if not q:
36 q = str(self.num2)
37 self.num2 -= 1
38 A.add_node(q, style = 'invis')
39 A.add_edge(p, q, style = 'invis')
40 if not r:
41 r = str(self.num2)
42 self.num2 -= 1
43 A.add_node(r, style = 'invis')
44 A.add_edge(p, r, style = 'invis')
45 l = str(self.num2)
46 self.num2 -= 1
47 A.add_node(l, style = 'invis')
48 A.add_edge(p, l, style = 'invis')
49 B = A.add_subgraph([q, l, r], rank = 'same')
50 B.add_edge(q, l, style = 'invis')
51 B.add_edge(l, r, style = 'invis')
52
53 return p #return key root node
54
55 self.A = pgv.AGraph(directed=True,strict=True)
56 writeHelp(self.tree.root, self.A)
57 self.A.graph_attr['epsilon']='0.001'
58 print self.A.string() # print dot file to standard output
59 self.A.layout('dot') # layout with dot
60 self.A.draw(outfile) # write to file
61
62
63 if __name__ == '__main__':
64 tree = BinaryTree()
65 tree.createTree(-1)
66 tree.inorderTravel()
67 writer = treeWriter(tree)
68 if len(sys.argv) > 1:
69 outfile = sys.argv[1]
70 writer.write(outfile) #write result to outfile
71 else:
72 writer.write() #write result to tree.png
73
74