huffman-archiver/internal/huffman/forest.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
}