huffman.py

#!/usr/bin/python
# coding: utf-8

from collections import Counter
from operator import itemgetter
import urllib
from matplotlib import pyplot as plt

'''
===========
 functions
===========
'''

def biSearch(a, x):
    l = 0
    r = len(a) - 1
    while (l <= r):
        mid = (l+r)/2
        if (a[mid][1] < x):
            r = mid - 1
        else:
            l = mid + 1
    return l

def generateTree(a):
    while len(a) > 1:
        l = a.pop()
        r = a.pop()
        val = l[1]+r[1]
        # node's structure is: [letter, value, left hand child, right hand child]
        node = [None, val, l, r]
        pos = biSearch(a, val) if len(a)>1 else 0
        a.insert(pos, node)

def uniqueGenerator():
    ''' the generator is needed to prevent the existance of nodes which are having more than one parent and more than two children
        normally it would occur if two nodes had the same value '''
    i = 0
    while True:
        yield ('i' + str(i))
        i += 1

def huffCode(tree, result, dot, values, code=[], dotTemp=[], isRight=False):
    if (tree[0] is not None):
        result[tree[0]] = ''.join(code), tree[1]
        dot.append(''.join(dotTemp) + tree[0] + ('', '[label=1]')[isRight] + ';\n')
    else:
        id = next(generator)
        values[id]=tree[1]
        huffCode(tree[2], result, dot, values, code+['0'], dotTemp + [id + '->'] if not isRight else [id + '->'], False)
        huffCode(tree[3], result, dot, values, code+['1'], [id + '->'] if not isRight else dotTemp + [id + '->'], True)

def generateDotScript(dot, huff, values):
    ''' http://en.wikipedia.org/wiki/DOT_(graph_description_language) '''
    letters = [ '{0}[shape=record,label="{{{{{0}|{1}}}|{2}}}"];\n'.format(k, v[1], v[0]) for k,v in huff.items() ]
    l = [ '{0}[label={1}];\n'.format(k, v) for k,v in values.items() ]
    return 'digraph G{\nedge[label=0];\ngraph[ranksep=0];\n' + ''.join(letters) + ''.join(l) + ''.join(dot) + '}'

def openDiagramImage(dot):
    ''' sends a POST request to Google Chart API '''
    try:
        dot = dot.encode('utf-8')
        data = urllib.urlencode({ 'cht': 'gv:dot',
                                  'chl': dot
        })
        url = 'http://chart.apis.google.com/chart'
        f = urllib.urlopen(url, data)
        a = plt.imread(f)
        plt.imshow(a)
        plt.axis('off')
        mng = plt.get_current_fig_manager()
        mng.resize(*mng.window.maxsize())
        plt.show()
    except:
        print "Couldn't open the image of generated tree. There might be a problem with your internet connection or the request to Google Chart API is too long."

'''
======
 main
======
'''

text = 'Litwo! Ojczyzno moja! ty jestes jak zdrowie. Ile cie trzeba cenic, ten tylko sie dowie, Kto cie stracil. Dzis pieknosc twa w calej ozdobie Widze i opisuje, bo tesknie po tobie. Panno Świeta, co jasnej bronisz Czestochowy I w Ostrej swiecisz Bramie! Ty, co grod zamkowy Nowogrodzki ochraniasz z jego wiernym ludem! Jak mnie dziecko do zdrowia powrocilas cudem (Gdy od placzacej matki pod Twoja opieke Ofiarowany, martwa podnioslem powieke I zaraz moglem pieszo do Twych swiatyn progu Isc za wrocone życie podziekowac Bogu), Tak nas powrocisz cudem na Ojczyzny lono. Tymczasem przenos moje dusze uteskniona Do tych pagorkow lesnych, do tych lak zielonych, Szeroko nad blekitnym Niemnem rozciagnionych; Do tych pol malowanych zbożem rozmaitem, Wyzlacanych pszenica, posrebrzanych żytem; Gdzie bursztynowy swierzop, gryka jak snieg biala, Gdzie panienskim rumiencem dziecielina pala, A wszystko przepasane, jakby wstega, miedza Zielona, na niej z rzadka ciche grusze siedza. Środ takich pol przed laty, nad brzegiem ruczaju, Na pagorku niewielkim, we brzozowym gaju, Stal dwor szlachecki, z drzewa, lecz podmurowany; Świecily sie z daleka pobielane sciany, Tym bielsze, że odbite od ciemnej zieleni Topoli, co go bronia od wiatrow jesieni. Dom mieszkalny niewielki, lecz zewszad chedogi, I stodole mial wielka, i przy niej trzy stogi Użatku, co pod strzecha zmiescic sie nie może; Widac, że okolica obfita we zboże, I widac z liczby kopic, co wzdluż i wszerz smugow Świeca gesto jak gwiazdy, widac z liczby plugow Orz cych wcze a snie lany ogromne ugoru, Czarnoziemne, zapewne należne do dworu, Uprawne dobrze na ksztalt ogrodowych grzadek: Że w tym domu dostatek mieszka i porzadek. Brama na wciaż otwarta przechodniom oglasza, Że goscinna i wszystkich w goscine zaprasza.'

# creates descendingly sorted list of tuples: (letter, frequency)
text = ''.join( i for i in text if i.isalpha())
text = text.upper()
letters = Counter(text)
letters = sorted(letters.iteritems(), key=itemgetter(1), reverse=True)

# generates a binary tree represented by a multilevel list
tree=letters
generateTree(tree)

generator = uniqueGenerator()
huff = {}
dot = []
values = {}
huffCode(tree[0], huff, dot, values)
for k,v in sorted(huff.iteritems(), key=itemgetter(0)):
    print '{0}: code={1}, value={2}'.format(k, v[0], v[1])

dotScript = generateDotScript(dot, huff, values)
openDiagramImage(dotScript)