84 lines
1.5 KiB
Go
84 lines
1.5 KiB
Go
|
package huffman
|
||
|
|
||
|
import (
|
||
|
"huffman/internal/io"
|
||
|
)
|
||
|
|
||
|
type TreeNode struct {
|
||
|
Value byte
|
||
|
Size int16
|
||
|
Frequency uint64
|
||
|
Left *TreeNode
|
||
|
Right *TreeNode
|
||
|
Parent *TreeNode
|
||
|
}
|
||
|
|
||
|
func (n *TreeNode) IsTerminate() bool {
|
||
|
return n.Left == nil && n.Right == nil
|
||
|
}
|
||
|
|
||
|
func (n *TreeNode) CalcSize() {
|
||
|
n.Size = 1
|
||
|
|
||
|
if n.Left != nil {
|
||
|
n.Left.CalcSize()
|
||
|
n.Size += n.Left.Size
|
||
|
}
|
||
|
|
||
|
if n.Right != nil {
|
||
|
n.Right.CalcSize()
|
||
|
n.Size += n.Right.Size
|
||
|
}
|
||
|
}
|
||
|
|
||
|
func DecodeTree(reader *io.BufferedReadSeeker) *TreeNode {
|
||
|
var stack []*TreeNode
|
||
|
|
||
|
for {
|
||
|
node := &TreeNode{}
|
||
|
|
||
|
if len(stack) > 0 {
|
||
|
node.Parent = stack[len(stack)-1]
|
||
|
|
||
|
if node.Parent.Left == nil {
|
||
|
node.Parent.Left = node
|
||
|
} else {
|
||
|
node.Parent.Right = node
|
||
|
}
|
||
|
}
|
||
|
|
||
|
stack = append(stack, node)
|
||
|
if reader.ReadBit() == 0x0 {
|
||
|
node.Value = reader.ReadByte()
|
||
|
|
||
|
for len(stack) > 1 && (stack[len(stack)-1].Right != nil || stack[len(stack)-1].IsTerminate()) {
|
||
|
stack = stack[:len(stack)-1]
|
||
|
}
|
||
|
|
||
|
if len(stack) == 1 && (stack[len(stack)-1].Right != nil || stack[len(stack)-1].IsTerminate()) {
|
||
|
break
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
stack[0].CalcSize()
|
||
|
return stack[0]
|
||
|
}
|
||
|
|
||
|
func encodeTreeRecursive(writer *io.BufferedWriter, root *TreeNode) {
|
||
|
if root.IsTerminate() {
|
||
|
writer.WriteBit(0x0)
|
||
|
writer.WriteByte(root.Value)
|
||
|
return
|
||
|
}
|
||
|
|
||
|
writer.WriteBit(0x1)
|
||
|
|
||
|
encodeTreeRecursive(writer, root.Left)
|
||
|
encodeTreeRecursive(writer, root.Right)
|
||
|
}
|
||
|
|
||
|
func EncodeTree(writer *io.BufferedWriter, root *TreeNode) {
|
||
|
encodeTreeRecursive(writer, root)
|
||
|
}
|