我正在使用 PyGraphviz 绘制二叉搜索树。我无法使用 PyGraphviz 创建重复节点,因此边缘循环回节点。
例如,以下代码仅生成 5 个节点,省略了重复节点。我尝试用唯一索引标记每个节点,但这并不能解决问题。
import pygraphviz as pgv
tree = pgv.AGraph(directed=True, strict=True)
tree.add_node(2)
tree.add_node(3)
tree.add_node(1)
tree.add_node(7)
tree.add_node(3)
tree.add_node(9)
tree.add_node(2)
tree.write('foo.dot')
image = pgv.AGraph('foo.dot')
image.layout()
image.draw('foo.pdf')
image.close()
我绘制到 BST 的代码:
import pygraphviz as pgv
import random
class Node:
insertion_step = []
def __init__(self, data):
self.data = data
self.left = None
self.right = None
def addNode(self, data):
if data < self.data:
if self.left is None:
self.left = Node(data)
self.printSubtree()
else:
self.left.addNode(data) # recursively calling addNode method
else:
if self.right is None:
self.right = Node(data)
self.printSubtree()
else:
self.right.addNode(data)
def printSubtree(self):
if not (self.left is None or self.right is None):
print self.left.data, self.data, self.right.data
self.insertion_step.append((self.left.data, self.data, self.right.data))
elif self.left is None and not self.right is None:
print None, self.data, self.right.data
self.insertion_step.append((None, self.data, self.right.data))
elif not self.left is None and self.right is None:
print self.left.data, self.data, None
self.insertion_step.append((self.left.data, self.data, None))
else:
print None, self.data, None
self.insertion_step.append((None, self.data, None))
def drawTree(self, tree, f):
print self.insertion_step
for step in self.insertion_step:
if not step[0] is None:
tree.add_node(step[0], color='goldenrod2', style='filled')
tree.add_node(step[1], color='goldenrod2', style='filled')
if not step[2] is None:
tree.add_node(step[2], color='goldenrod2', style='filled')
if step[0] is None or step[1] is None or step[2] is None:
tree.add_node('', color='goldenrod1', shape='box', style='filled')
if not step[0] is None:
tree.add_edge(step[1], step[0], color='sienna', style='filled')
else:
tree.add_edge(step[1], '', color='sienna', style='filled')
if not step[2] is None:
tree.add_edge(step[1], step[2], color='sienna', style='filled')
else:
tree.add_edge(step[1], '', color='sienna', style='filled')
tree.write(f)
img = pgv.AGraph(f)
img.layout()
img.draw(f.split('.')[0] + '.pdf')
img.close()
if __name__ == '__main__':
lst = [random.randint(1, 10) for i in range(10)]
print lst
n = Node(lst[0])
n.printSubtree()
for num in lst[1:]:
n.addNode(num)
tree = pgv.AGraph(directed=True, strict=True)
filename = 'tree.dot'
n.drawTree(tree, filename)
从上图中可以看出,边是循环的,因为没有创建重复节点。请建议我实现这一目标的方法。图中方框代表空节点。