1
\$\begingroup\$

Here is my implementation of Huffman encoding using min-heap approach mentioned in this Wikipedia page. I'm looking for advice, using Python 2.7.

import heapq
from collections import defaultdict
class TreeNode:
    def __init__(self, name, left, right, value):
        self.name = name
        self.left = left
        self.right = right
        self.value = value
    def __cmp__(self, other):
        return cmp(self.value, other.value)
    @staticmethod
    def merge_tree(root1, root2):
        return TreeNode(root1.name+"_"+root2.name, root1, root2, root1.value + root2.value)
    @staticmethod
    def print_tree(root):
        if not root:
            return
        buf = []
        buf.append(root)
        layer_count = 1
        next_layer_count = 0
        layer_result = []
        while len(buf) > 0:
            node = buf.pop(0)
            layer_result.append(str(node.name))
            layer_count -= 1
            if node.left:
                buf.append(node.left)
                next_layer_count += 1
            if node.right:
                buf.append(node.right)
                next_layer_count += 1
            if layer_count == 0:
                print '\t'.join(layer_result)
                layer_result = []
                layer_count = next_layer_count
                next_layer_count = 0
    @staticmethod
    def encode(root, prefix, result):
        if not root:
            return
        # left encode 0, right encode 1
        if not root.left and not root.right:
            result[root.name] = prefix
        TreeNode.encode(root.left, prefix+'0', result)
        TreeNode.encode(root.right, prefix+'1', result)

if __name__ == "__main__":
    n1 = TreeNode('a1', None, None, 0.4)
    n2 = TreeNode('a2', None, None, 0.35)
    n3 = TreeNode('a3', None, None, 0.2)
    n4 = TreeNode('a4', None, None, 0.05)
    h = [n1, n2, n3, n4]
    heapq.heapify(h)
    while len(h) > 1:
        c1 = heapq.heappop(h)
        c2 = heapq.heappop(h)
        heapq.heappush(h, TreeNode.merge_tree(c1, c2))

    TreeNode.print_tree(h[0])
    result = defaultdict(str)
    TreeNode.encode(h[0],'', result)
    print result
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.