90 lines
1.6 KiB
Go
90 lines
1.6 KiB
Go
|
package huffman
|
||
|
|
||
|
import (
|
||
|
. "huffman/internal/io"
|
||
|
"io"
|
||
|
"math"
|
||
|
)
|
||
|
|
||
|
func encodeNode(root, prev *TreeNode) (res []byte) {
|
||
|
if root.Parent != nil {
|
||
|
res = encodeNode(root.Parent, root)
|
||
|
}
|
||
|
|
||
|
if prev == nil {
|
||
|
return
|
||
|
}
|
||
|
|
||
|
if root.Left == prev {
|
||
|
res = append(res, 0x0)
|
||
|
} else {
|
||
|
res = append(res, 0x1)
|
||
|
}
|
||
|
|
||
|
return
|
||
|
}
|
||
|
|
||
|
func initEncodeMap(root *TreeNode) map[byte][]byte {
|
||
|
m := map[byte][]byte{}
|
||
|
q := []*TreeNode{root}
|
||
|
|
||
|
for len(q) > 0 {
|
||
|
cur := q[0]
|
||
|
q = q[1:]
|
||
|
|
||
|
if cur.Left != nil {
|
||
|
q = append(q, cur.Left)
|
||
|
}
|
||
|
|
||
|
if cur.Right != nil {
|
||
|
q = append(q, cur.Right)
|
||
|
}
|
||
|
|
||
|
if cur.IsTerminate() {
|
||
|
m[cur.Value] = encodeNode(cur, nil)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
return m
|
||
|
}
|
||
|
|
||
|
func calculateContentLen(root *TreeNode, len uint64) uint64 {
|
||
|
if root.IsTerminate() {
|
||
|
return len * root.Frequency
|
||
|
}
|
||
|
|
||
|
return calculateContentLen(root.Left, len+1) + calculateContentLen(root.Right, len+1)
|
||
|
}
|
||
|
|
||
|
func Encode(reader io.ReadSeeker, writer io.Writer) {
|
||
|
bufReader := NewBufferedReader(reader)
|
||
|
bufWriter := NewBufferedWriter(writer)
|
||
|
|
||
|
tree, size := CreateTree(bufReader)
|
||
|
treeMap := initEncodeMap(tree)
|
||
|
lg := int(math.Ceil(math.Log2(float64(size)))) + 7
|
||
|
|
||
|
EncodeTree(bufWriter, tree)
|
||
|
|
||
|
if tree.Size == 1 {
|
||
|
bufWriter.WriteByte(byte(lg))
|
||
|
bufWriter.WriteBitsFromNumber(size, lg)
|
||
|
} else {
|
||
|
treeBitLen := uint64(5*tree.Size + 4)
|
||
|
contentBitLen := calculateContentLen(tree, 0)
|
||
|
|
||
|
additionalBits := (8 - (contentBitLen+treeBitLen+3)%8) % 8
|
||
|
|
||
|
bufWriter.WriteBitsFromNumber(additionalBits, 3)
|
||
|
}
|
||
|
|
||
|
bufReader.SeekFromStart(0)
|
||
|
|
||
|
for !bufReader.IsEOF() {
|
||
|
node := treeMap[bufReader.ReadByte()]
|
||
|
bufWriter.WriteBits(node)
|
||
|
}
|
||
|
|
||
|
bufWriter.Flush()
|
||
|
}
|