python networkx pygraphviz简单使用

from heapq import *
import networkx as nx
import matplotlib.pyplot as plt
import pygraphviz as pgv

class Node:

    def __init__(self, label, weight):
        self.left = None
        self.right = None
        self.label = label
        self.weight = weight
    def __lt__(self, other):
        return self.weight < other.weight

    def __eq__(self, other):
        return self.weight == other.weight and self.label==other.label

    def __hash__(self):
        return hash(id(self))
    def __str__(self):
        return str(self.weight) if not self.label.startswith('N') else self.label+':'+str(self.weight)


def huffman(A):
    A=[Node(str(i),i)for i in A]
    cnt=1
    heapify(A)
    while len(A)>1:
            x=heappop(A);y=heappop(A)
            xy=Node('N'+str(cnt),x.weight+y.weight)
            xy.left=x;xy.right=y

            heappush(A,xy);cnt+=1
            if len(A)==1:
                A[0].label='Root'
                return heappop(A)

def mplt(A):
    T=nx.DiGraph()
    if  A is not None:
        cons(A,T)
    nx.draw(T, with_labels=True, font_weight='bold')

    G=pgv.AGraph(strict=True,directed=True)
    G.add_nodes_from(T.nodes())
    G.add_edges_from(T.edges())

    G.write("first.dot")


    G.layout(prog='dot')

    G.draw('/Users/hello/first.png')
    plt.show()

def cons(A,T):
    if  A is not None:
        if A.left is not None:
            T.add_edge(A,A.left)
            cons(A.left,T)
        if A.right is not None:
            T.add_edge(A,A.right)
            cons(A.right,T)

def visit(A):
    if  A is not None:
        print(A.label ,A.weight)
        visit(A.left)
        visit(A.right)
if __name__ == '__main__':
        A=[1,45,4,2,3,6,7,8,9,22,111,231]
        h_g=huffman(A)
        visit(h_g)
        mplt(h_g)