huff/huffman.go

66 lines
1.3 KiB
Go

package huff
import (
"container/heap"
)
type Node struct {
Weight int64
Value interface{}
Left *Node
Right *Node
}
func (n Node) Lookup(value interface{}) *Node {
return nil
}
func NewWordTree(words []string) *Node {
nodes := []*Node{}
for i := range words {
nodes = append(nodes, &Node{Weight: int64(i + 1), Value: words[i]})
}
h := nodeInterface(nodes)
h_ptr := &h
heap.Init(h_ptr)
count := h_ptr.Len()
for count > 1 {
left := heap.Pop(h_ptr)
right := heap.Pop(h_ptr)
heap.Push(h_ptr, Combine(left.(*Node), right.(*Node)))
count = h_ptr.Len()
}
return heap.Pop(h_ptr).(*Node)
}
func Combine(left *Node, right *Node) *Node {
var weight int64
if left != nil {
weight = weight + left.Weight
}
if right != nil {
weight = weight + right.Weight
}
return &Node{
Weight: weight,
Left: left,
Right: right,
}
}
type nodeInterface []*Node
func (ni nodeInterface) Len() int { return len(ni) }
func (ni nodeInterface) Less(i, j int) bool { return ni[i].Weight < ni[j].Weight }
func (ni nodeInterface) Swap(i, j int) { ni[i], ni[j] = ni[j], ni[i] }
func (ni *nodeInterface) Push(x interface{}) { *ni = append(*ni, x.(*Node)) }
func (ni *nodeInterface) Pop() interface{} {
old := *ni
n := len(old)
x := old[n-1]
*ni = old[0 : n-1]
return x
}