python
from collections import Counter
import heapq
from itertools import count
def build_huffman_codes(text):
if not text:
return {}
frequencies = Counter(text)
serial = count()
heap = []
for char, freq in frequencies.items():
heapq.heappush(
heap,
(freq, next(serial), char)
)
if len(heap) == 1:
only_char = next(iter(frequencies))
return {only_char: "0"}
while len(heap) > 1:
freq1, _, left = heapq.heappop(heap)
freq2, _, right = heapq.heappop(heap)
parent = (left, right)
heapq.heappush(
heap,
(freq1 + freq2, next(serial), parent)
)
root = heap[0][2]
codes = {}
def walk(node, prefix):
if isinstance(node, str):
codes[node] = prefix
return
left, right = node
walk(left, prefix + "0")
walk(right, prefix + "1")
walk(root, "")
return codes
text = "AAAAABBC"
codes = build_huffman_codes(text)
encoded = "".join(codes[c] for c in text)
print(codes)
print(encoded)