来自二叉树的Edgelist

时间:2012-11-22 18:48:54

标签: python binary-tree networkx

我想制作一张漂亮的二叉树图。

这是我的自定义BinaryTree类:

class BinaryTree():

   def __init__(self, data):
      self.data = data
      self.right = None
      self.left = None

现在,为了绘制这个图,我将使用networkx库,所以我需要将我的图转换为networkx对象,然后使用graphviz绘制它。问题是边缘列表:为了构建我的新对象,我需要边缘。

例如,给出如下图所示的二叉树。 enter image description here

我需要检索边缘列表。会是这样的:

[(0,1),(0,2),(2,3),(2,4)]

请注意,在我的情况下,我没有节点上的id。那么我怎么能这样做呢? 我相信它可能是一些考虑到深度的递归函数,但我遇到了一些困难,所以有点帮助是值得赞赏的。 ;)

修改

感谢您的回答。但我发现自己的解决方案效果很好..:P 这是:

def edgelist(node, output, id=0):

    if node is None or isinstance(node, bt.Leaf):
         return output

    if node.left:
         output.append((id, id*2+1))

    if node.right:
         output.append((id, id*2+2))

    edgelist(node.left, output, id*2+1)
    edgelist(node.right, output, id*2+2)

    return output

2 个答案:

答案 0 :(得分:1)

以下是修改BinaryTree类以转储边缘列表的一种方法:

import networkx as nx
import itertools as IT
import matplotlib.pyplot as plt

class BinaryTree(object):
   def __init__(self, data):
      self.data = data
      self.right = None
      self.left = None
      self.name = None
   def edgelist(self, counter = IT.count().next):
       self.name = counter() if self.name is None else self.name
       for node in (self.left, self.right):       
           if node:
               node.name = counter() if node.name is None else node.name
               yield (self.name, node.name)
       for node in (self.left, self.right):
           if node:
               for n in node.edgelist(counter):
                   yield n

tree = [BinaryTree(i) for i in range(5)]        
tree[0].left = tree[1]
tree[0].right = tree[2]
tree[2].left = tree[3]
tree[2].right = tree[4]

edgelist = list(tree[0].edgelist())
print(edgelist)   

G = nx.Graph(edgelist)
nx.draw_spectral(G)
plt.show()

产量

[(0, 1), (0, 2), (2, 3), (2, 4)]

enter image description here

答案 1 :(得分:0)

您可以使用collections.dequeue来避免递归:

import collections
def edges_breadth(tree):
    history = collections.deque([tree])
    while history:
        parent = history.popleft()
        for c in (parent.left, parent.right):
            if c:
                yield((parent.data, c.data))
                history.append(c)

请注意,这是广度优先遍历。您可能需要另一个traversal order,即深度优先,如此预订递归实现:

def edges_depth(tree):
    results = []
    def visit(parent, child):
        if child:
            results.append((parent, child))
            visit(child.left)
            visit(child.right)
    return results