1package huffman
  2
  3import (
  4	"fmt"
  5	"sort"
  6)
  7
  8const (
  9	LEAF = 0
 10	MID  = 1
 11)
 12
 13type Node struct {
 14	typ  int64
 15	key  int64
 16	leaf byte
 17	l    *Node
 18	r    *Node
 19}
 20
 21func (n *Node) String() string {
 22	if n.typ == LEAF {
 23		return fmt.Sprintf("%s (%v)", string(n.leaf), n.key)
 24	}
 25	return fmt.Sprintf("%v", n.key)
 26}
 27
 28func newLeaf(b byte, k int64) *Node {
 29	return &Node{
 30		typ:  LEAF,
 31		key:  k,
 32		leaf: b,
 33	}
 34}
 35
 36func newNode(a, b *Node) *Node {
 37	var l, r *Node
 38	if a.key < b.key {
 39		l = a
 40		r = b
 41	} else {
 42		l = b
 43		r = a
 44	}
 45	n := &Node{
 46		typ: MID,
 47		key: a.key + b.key,
 48		l:   l,
 49		r:   r,
 50	}
 51	return n
 52}
 53
 54func leafsFromValues(values map[byte]int64) []*Node {
 55	var nodes []*Node
 56	for k, v := range values {
 57		n := newLeaf(k, v)
 58		nodes = append(nodes, n)
 59	}
 60	return nodes
 61}
 62
 63func rmFromNodes(i int, nodes []*Node) []*Node {
 64	if i == len(nodes) {
 65		return nodes[:i]
 66	}
 67	return append(nodes[:i], nodes[i+1:]...)
 68}
 69
 70// sortNodes returns a []*Node sorted
 71func sortNodes(nodes []*Node) []*Node {
 72	sort.Slice(nodes, func(i, j int) bool {
 73		return nodes[i].key < nodes[j].key
 74	})
 75	return nodes
 76}
 77
 78func buildTree(nodes []*Node) *Node {
 79	// build binary tree
 80	for len(nodes) > 1 {
 81		// pair two smallest nodes
 82		nodes = sortNodes(nodes)
 83		// create a new node with them
 84		n := newNode(nodes[0], nodes[1])
 85		// remove the picked nodes from the array of nodes
 86		nodes = rmFromNodes(0, nodes)
 87		nodes = rmFromNodes(0, nodes)
 88		// add the new node (the two old paired) to the node array
 89		nodes = append(nodes, n)
 90	}
 91	return nodes[0]
 92}
 93
 94func descendTable(n *Node, path string, table map[byte]string) map[byte]string {
 95	if n.typ == MID {
 96		table = descendTable(n.l, path+"0", table)
 97		table = descendTable(n.r, path+"1", table)
 98	} else if n.typ == LEAF {
 99		table[n.leaf] = path
100	}
101	return table
102}
103
104func generateTable(n *Node) map[byte]string {
105	table := make(map[byte]string)
106	table = descendTable(n, "", table)
107	return table
108}
109
110func Huffman(b []byte) error {
111	// compute frequencies
112	values := make(map[byte]int64)
113	for _, v := range b {
114		values[v]++
115	}
116	leafs := leafsFromValues(values)
117
118	// sort frequencies
119	sortedLeafs := sortNodes(leafs)
120
121	// build binary tree
122	n := buildTree(sortedLeafs)
123	fmt.Println(n)
124
125	// get the table (binary code for each value)
126	table := generateTable(n)
127	fmt.Println(table)
128
129	// WIP
130
131	return nil
132}