huffman2.py 2.08 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
#!/usr/bin/env python3

import subprocess

class Noeud(object):
    def __init__(self,p=0,left=None,right=None,code='',name=''):
        self.left = left
        self.right= right
        if left != None and right != None :
            self.p = left.p + right.p
            self.name =left.name+right.name
        else:
            self.p = p
            self.name = name
        self.code = code
    def __lt__(self,other):
        return self.p<other.p
    def __repr__(self):
        return self.name

table = [
    ('A', 25),
    ('B', 20),
    ('C', 15),
    ('D', 12),
    ('E', 10),
    ('F', 8),
    ('G', 5),
    ('H', 5)]

table_noeud = [Noeud(name=x[0],p=x[1]) for x in table]
#print(table_noeud)
def create_tree(table_noeud):
    queue = table_noeud.copy()
    while len(queue) > 2:
        queue.sort()
        print(queue)
        l=queue.pop(0)
        r=queue.pop(0)
        queue.append(Noeud(left=l,right=r))
    root= Noeud(left=queue[0],right=queue[1])
    return root
root_node= create_tree(table_noeud)
print(root_node)

def gen_code(node,prefix=''):
    if node.left != None:
        node.code = prefix
        gen_code(node.left,prefix+'0')
        gen_code(node.right,prefix+'1')
    else:
        node.code = prefix
        print(node.name,node.code)

gen_code(root_node)

def draw_tree(node):
    if len(node.name) == 1: # feuille
        desc ='N{} [label="{}:{}", fontcolor=blue, fontsize=16, width=2, shape=box];\n'.format(node.code, node.name, node.code)
    else:
        desc = 'N{} [label="{}"];\n'.format(node.code,node.code)
        desc += draw_tree(node.left)
        desc += draw_tree(node.right)
        desc += 'N{}:n -> N{}:e;\n'.format(node.code,node.left.code)
        desc += 'N{}:s -> N{}:e;\n'.format(node.code,node.right.code)
    return desc

with open('graph.dot','w') as f:
    f.write('digraph G {\n ')
    f.write(' splines=ortho \n')
    f.write('rankdir=RL;\n')
    f.write(draw_tree(root_node))
    f.write('{rank =same; N' + '; N'.join([n.code for n in table_noeud]) +';}\n')
    f.write('}')
subprocess.call('dot -Tpng graph.dot -o graph.png', shell=True)
print('done')