Huffman.py

import networkx as nx
import matplotlib.pyplot as plot
from sys import argv

vertices = dict()

#Clase Nodo, estructura de informacion
class Nodo:
	def __init__(self, nombre, frecuencia, izquierda, derecha):
		self.nombre = nombre
		self.frecuencia = frecuencia
		self.izquierda = izquierda
		self.derecha = derecha
		self.codigo = ""

	def __cmp__( self, other ):
		if other != None:
			if self.frecuencia < other.frecuencia:  res = -1
			elif self.frecuencia > other.frecuencia:  res = 1
			else :  res = 0
			return res
		else:
			return 1


class Huffman():
	def __init__(self, n):
		self.n = n

	def ordenar(self, d):
		listaDeNodos = list()
		letras = d.items()
		letras.sort(key=lambda x:x[1])
		for e in letras:
			nodo = Nodo(e[0],e[1],None, None)
			listaDeNodos.append(nodo)
		return listaDeNodos

	def obtenerValorBinario(self, cadena):
		res = ''
		for e in cadena:
			res += bin(int(ord(e)))[2:]
		return res

	def obtenerFrecuencias(self, cadena):
		f = dict()
		for e in cadena:
			if f.has_key(e):
				f[e] += 1
			else:
				f[e] = 1
		return f

	def generaNuevoNodo(self, n1, n2):
		nombre = n1.nombre + n2.nombre
		suma = n1.frecuencia + n2.frecuencia
		n1.codigo = '0'
		n2.codigo = '1'
		nuevo = Nodo(nombre, suma, n1, n2)
		return nuevo

	def generaCodigos(self, r):
		if r.izquierda != None:
			r.izquierda.codigo = r.codigo + r.izquierda.codigo
			self.generaCodigos(r.izquierda)
		if r.derecha != None:
			r.derecha.codigo = r.codigo + r.derecha.codigo
			self.generaCodigos(r.derecha)

	def stl(self, c):
		l = list()
		for e in c:
			l.append(e)
		return l

	def pintaGrafo(self, elementos):
		print "nodo:",r.nombre
		print "valor ", r.frecuencia
		print "codigo ", r.codigo
		print ""
		if not elementos.has_key(r.nombre):
			elementos[r.nombre] = r
		if r.izquierda != None:
			elementos = pintaGrafo(r.izquierda, elementos)
		if r.derecha != None:
			elementos = pintaGrafo(r.derecha, elementos) 
		return elementos

	def obtenerDiccNodos(self, r, elementos):
		if not elementos.has_key(r.nombre):
			elementos[r.nombre] = r
		if r.izquierda != None:
			elementos = self.obtenerDiccNodos(r.izquierda, elementos)
		if r.derecha != None:
			elementos = self.obtenerDiccNodos(r.derecha, elementos) 
		return elementos

	def obtenerMapaDeCodigos(self, d):
		nombreYCodigo = dict()
		llaves = d.keys()
		for e in llaves:
			nodo = d[e]
			nombreYCodigo[nodo.nombre] = nodo.codigo
		return nombreYCodigo

	#Funcion que genera el codigo binario en base al texto original y el banco de codigos
	def generaCodigoHuffman(self, texto, codigos):
		huffman = ''
		for e in texto:
			huffman += codigos[e]
		return huffman

	#funcion que descomprime el codigo huffman en base al arbol final
	def descomprime(self, arbol, codigo):
		original = arbol
		lcodigo = self.stl(codigo)
		msj = ''
		bandera = True
		while (len(lcodigo) > 0):
			v = int(lcodigo.pop(0))
			if(v == 0):
				if (arbol.izquierda != None and len(arbol.izquierda.nombre) == 1):
					msj += arbol.izquierda.nombre
					arbol = original
				else:
					arbol = arbol.izquierda
			else:
				if (arbol.derecha != None and len(arbol.derecha.nombre) == 1):
					msj += arbol.derecha.nombre
					arbol = original
				else:
					arbol = arbol.derecha
		return msj

	#Genera diccionario con estructura {(vertice1,vertice2):valorDeArista}
	def obtenerVertices(self, nodo, padre):
		global vertices
		if padre != None:
			vertices[(padre.nombre, nodo.nombre)] = nodo.frecuencia
		if nodo.izquierda != None:
			self.obtenerVertices(nodo.izquierda, nodo)
		if nodo.derecha != None:
			self.obtenerVertices(nodo.derecha, nodo)
		return vertices

	#Funcion para pintar el arbol con networkx
	def generarArbol(self, vertices):
		G=nx.MultiDiGraph()
		listaVertices = vertices.keys()
		for llave in listaVertices:
			vert = vertices[llave]
			G.add_edge(llave[0],llave[1])
		pos=nx.spring_layout(G)
		nx.draw_networkx_nodes(G,pos, node_color="white")
		nx.draw_networkx_edges(G,pos, edge_color='blue')
		nx.draw_networkx_labels(G,pos)
		nx.draw_networkx_edge_labels(G,pos, vertices)
		plot.axis('off')
		plot.show()


	def proceso(self, cadena):
		cadbin = self.obtenerValorBinario(cadena)
		frec = self.obtenerFrecuencias(cadena)
		lista = self.ordenar(frec)
		while (len(lista) > 1):
			nuevo = self.generaNuevoNodo(lista.pop(0), lista.pop(0))
			lista.append(nuevo)
			lista.sort()
		ARBOL = lista[0]
		self.generaCodigos(ARBOL)
		elementos = dict()
		elementos = self.obtenerDiccNodos(ARBOL, elementos)
		codigos = self.obtenerMapaDeCodigos(elementos)
		huffman = self.generaCodigoHuffman(cadena, codigos)

		return huffman, ARBOL