78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
package huffman
|
|
|
|
import "huffman/internal/io"
|
|
|
|
type ForestNode struct {
|
|
Frequency uint64
|
|
Root *TreeNode
|
|
}
|
|
|
|
type Forest struct {
|
|
Nodes [256]ForestNode
|
|
Size int
|
|
}
|
|
|
|
func CreateForest(reader *io.BufferedReadSeeker) *Forest {
|
|
var forest Forest
|
|
forest.Size = 256
|
|
|
|
for !reader.IsEOF() {
|
|
forest.Nodes[reader.ReadByte()].Frequency++
|
|
}
|
|
|
|
for i := 0; i < forest.Size; i++ {
|
|
forest.Nodes[i].Root = &TreeNode{Size: 1, Value: byte(i), Frequency: forest.Nodes[i].Frequency}
|
|
}
|
|
|
|
return &forest
|
|
}
|
|
|
|
func (f *Forest) TwoMin() (idx1 int, idx2 int) {
|
|
idx2 = 1
|
|
|
|
if f.Nodes[idx1].Frequency > f.Nodes[idx2].Frequency {
|
|
idx1, idx2 = idx2, idx1
|
|
}
|
|
|
|
for i := 2; i < f.Size; i++ {
|
|
if f.Nodes[idx1].Frequency > f.Nodes[i].Frequency {
|
|
idx1, idx2 = i, idx1
|
|
} else if f.Nodes[idx2].Frequency > f.Nodes[i].Frequency {
|
|
idx2 = i
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
func CreateTree(reader *io.BufferedReadSeeker) (*TreeNode, uint64) {
|
|
f := CreateForest(reader)
|
|
|
|
for f.Size > 1 {
|
|
idx1, idx2 := f.TwoMin()
|
|
node1, node2 := f.Nodes[idx1], f.Nodes[idx2]
|
|
|
|
f.Nodes[idx1].Root = &TreeNode{
|
|
Size: 1 + node1.Root.Size + node2.Root.Size,
|
|
Left: node1.Root,
|
|
Right: node2.Root,
|
|
}
|
|
|
|
if node1.Frequency == 0 {
|
|
f.Nodes[idx1].Root = node2.Root
|
|
} else if node2.Frequency == 0 {
|
|
f.Nodes[idx1].Root = node1.Root
|
|
}
|
|
|
|
f.Nodes[idx1].Frequency += node2.Frequency
|
|
if node1.Frequency != 0 && node2.Frequency != 0 {
|
|
node1.Root.Parent = f.Nodes[idx1].Root
|
|
node2.Root.Parent = f.Nodes[idx1].Root
|
|
}
|
|
|
|
f.Nodes[idx2] = f.Nodes[f.Size-1]
|
|
f.Size--
|
|
}
|
|
|
|
return f.Nodes[0].Root, f.Nodes[0].Frequency
|
|
}
|