利用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

简单定义了二叉链表节点，和二叉树类，给出了通过用户交互输入生成二叉树的方法，及中序遍历方法。

1  # Notice the differnce of creating tree code with c++, python does not
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 文件并进一步生成二叉树图案输出到文件。

绘制过程采用递归，深度优先遍历，利用了不可见节点和边从而能够始终确保能够分清是左子树还是又子树，如果不采用

加入不可见节点，边的方法则无法分清，如果绘制树的话可以采用简单的不加不可见节点的方法。

1  import sys
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