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}