golang哈夫曼编码压缩文件代码实现全流程(超详细版)

一、最近研究了一些通用的压缩算法,发现目前各大博客中和相关教程中关于使用golang实现哈夫曼压缩算法的好文章屈指可数,大多是实验性的代码,并没有完全实现压缩文件的所有必要步骤

1.仅仅介绍了哈夫曼树的机制,并没有算法实现,只知道原理可不一定能写出代码
2.代码实现了哈夫曼树的构造,并且完成编码,存储在string变量中,打印出被哈夫曼编码过的字符串,不知道这样做的意义何在,这其实变相增大了文件的体积,是一种耍流氓行为,压根没有达到目的
3.通过全局变量将编码的table保存在内存中,而非保存在被压缩的文件中,这也是一种耍流氓行为,难道我们需要拿到原文件和压缩文件才能解压?
本文较长,想了解某个模块实现的同学可以跳转到标题四后如下序列
1.词频统计,2. 哈夫曼树的实现,3.从哈夫曼树得到编码表,4.将经过哈夫曼编码的字符写入文件中
5.将编码表写入文件头的细节,6-7.解压缩文件的细节

二、为了解决以上所述的痛点,真正用golang写出能够压缩解压缩文件的代码,我将着重解决以下几个问题

  1. 结合golang的特性,展示构建哈夫曼树过程中heap接口的应用,用足够简单的代码构建哈夫曼树
  2. 用位运算将被编码的字符写入bit中,真正实现对字符集的压缩
  3. 用io操作,展现读写文件头table的细节,真正保证压缩文件具有全量的信息

三、Talk is cheap show me the code

先贴出全量代码 , 文件输入输出路径均定义在代码开头的const变量中,有想验证代码可用性的同学可以自行拷贝实验

package main

import (
    "bufio"
    "bytes"
    "container/heap"
    "crypto/md5"
    "encoding/hex"
    "fmt"
    "io/ioutil"
    "log"
    "os"
    "sort"
    "strconv"
    "strings"
)

const filePath = `D:\dataCourpus\file.010`
const outPutPath = `D:\dataCourpus\outPut`
const depressPath = `D:\dataCourpus\depress.txt`
const contentBuffer = 5000000


type TreeNode struct {
    Val int
    Times int
    Left *TreeNode
    Right *TreeNode
}

type  treeHeap []*TreeNode

func (p treeHeap) Less( i , j int) bool {
    return p[i].Times <= p[j].Times
}

func (p treeHeap) Len() int {
    return len(p)
}

func (p treeHeap) Swap(i , j int) {
    p[i] , p[j] = p[j] , p[i]
}

func (p *treeHeap) Push(node interface{}) {
    *p = append(*p, node.(*TreeNode))
}

func (p *treeHeap) Pop()  interface{} {
    n := len(*p)
    t := (*p)[n-1]
    *p = (*p)[:n-1]
    return t
}

// 测试压缩解压缩的正确性
func main() {
    HuffmanEncoding(filePath,outPutPath)
    depress(outPutPath,depressPath)
    originMD5 , _:= MD5File(filePath)
    recoverMD5 , _ := MD5File(depressPath)
    fmt.Println(originMD5 == recoverMD5)
}

func HuffmanEncoding (filePath , outPath string) {
    // 思路: 1. 读取文本内容,存放到内存中,或者以流的形式读取文本内容,构建二叉树即可。
    // 统计每个字出现的频次
    file ,err := os.Open(filePath)
    if err != nil  {
        log.Fatalln(err)
        return
    }
    defer file.Close()
    // 我们不需要关心总量是多少,因为分母是固定的,只需要知道频率按次数排序即可。
    imap := getFrequencyMap(file)
    plist := make(treeHeap,0)
    // 遍历map ,将键值对存入pair,然后按频率排序
    for k , v := range imap {
        plist = append(plist, &TreeNode{Val: k,Times: v})
    }
    sort.Sort(plist)
    //如果文件是空的,还构造个屁
    if len(plist) == 0 {return}
    hTree := initHuffmanTree(plist)
    /*遍历哈弗曼树,生成哈夫曼编码表(正表,用于编码),key(ASSCII),value(路径痕迹)*/
    encodeMap := make(map[int]string)
    createEncodingTable(hTree , encodeMap)

    // 将输入文件的字符通过码表编码,输出到另一个文件 , 压缩模块完成
    encoding(filePath , outPath , encodeMap)

}

func writeTable(path string, codeMap map[int]string , left int) {
    file ,err := os.Create(path)
    if err != nil {
        return
    }
    // 第一行,写入文件头的长度
    var buff bytes.Buffer
    buff.WriteString(strconv.Itoa(len(codeMap)+1) + "\n")
    for k , v := range codeMap {
        buff.WriteString(strconv.Itoa(k) + ":" + v + "\n")
    }
    buff.WriteString(strconv.Itoa(left) + "\n")
    file.WriteString(buff.String())
    file.Close()
}


/* 一次性读入,存到string或者buffer.string中 */
func encoding(inPath string, outPath string , encodeMap map[int]string) {
    /*1.先尝试一次性读入*/
    inFile ,err  := os.Open(inPath)
    defer inFile.Close()
    if err != nil {
        return
    }
    reader := bufio.NewReader(inFile)
    fileContent := make([]byte,contentBuffer)
    count , _ := reader.Read(fileContent)
    var buff bytes.Buffer
    //string编码
    for i := 0 ; i < count ; i ++  {
        v := fileContent[i]
        if code , ok := encodeMap[int(v)] ; len(code)!= 0  && ok  {
            buff.WriteString(code)
        }
    }
    res := make([]byte,0)
    var buf byte = 0
    //bit编码
    //TODO 记录bit剩余位,很简单只要对buff.bytes取长度对8取余即可
    for idx , bit := range buff.Bytes() {
        //每八个位使用一个byte读取,结果存入res数组即可
        pos := idx % 8
        if pos == 0 && idx > 0 {
            res = append(res, buf)
            buf = 0
        }
        if bit == '1' {
            buf  |=   1 <<  pos
        }
    }
    //TODO 这个left是剩余待处理的位数
    left := buff.Len() % 8
    res = append(res, buf)
    // 将编码数组写入文件 , TODO 先将码表和left数写入文件,解码时在开头读取
    writeTable(outPath,encodeMap,left)
    outFile , err := os.OpenFile(outPath,os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil{
        return
    }
    wcount , err := outFile.Write(res)
    if err != nil {
        fmt.Println(wcount)
        return
    }
}

//码表,考虑到性能必须要生成 map, key(int对应ASSCII🐎,string对应bit编码,后续转成bit)
func createEncodingTable(node *TreeNode , encodeMap map[int]string )  {
    /*思路:回溯遍历二叉树,用byte记录0 1🐎,遇到叶子节点就转换成string存入码表
        遍历顺序:根左右
    */
    tmp := make([]byte,0)
    var depth func (treeNode *TreeNode)
    depth = func (root *TreeNode) {
        //如果已经遍历到空,返回
        if root == nil {
            return
        }
        //如果遍历到的是叶子节点 , byte转换成string,入表
        if root.Left == nil && root.Right == nil {
            encodeMap[root.Val] = string(tmp)
        }
        //如果是普通节点,左右递归回溯即可 #规则: 左0右1
        tmp = append(tmp ,'0' )
        depth(root.Left)
        tmp[len(tmp)-1] = '1'
        depth(root.Right)
        tmp = tmp[:len(tmp)-1]
    }
    depth(node)
}

func initHuffmanTree(plist treeHeap) *TreeNode {
    //使用优先队列构造最小路径权值哈夫曼树
    heap.Init(&plist)
    for plist.Len() > 1 {
        t1 := heap.Pop(&plist).(*TreeNode)
        t2 := heap.Pop(&plist).(*TreeNode)
        root := &TreeNode{Times: t1.Times + t2.Times}
        if t1.Times > t2.Times {
            root.Right  , root.Left = t1 , t2
        }else {
            root.Right , root.Left = t2 , t1
        }
        heap.Push(&plist,root)
    }
    return plist[0]
}

func getFrequencyMap( file *os.File ) map[int]int {
    imap := make(map[int]int)
    // 读入文件数据,readline 记入map中,统计频次
    // 注意:Create不区分文件名大小写
    reader := bufio.NewReader(file)
    buffer := make([]byte,contentBuffer)
    readCount , _ := reader.Read(buffer)
    for i := 0 ; i < readCount ; i ++ {
        imap[int(buffer[i])] ++
    }
    return imap
}


func depress( inPath , depressPath string) {
    // originPath 原文件(或者可以传入码表), inPath 读入被压缩的文件 , depressPath 还原后的输出路径
    encodeMap := make(map[int]string)
    decodeMap := make(map[string]int)
    //2.读入压缩文件
    compressFile , _ := os.Open(inPath)
    // br 读取文件头 ,返回偏移量
    br := bufio.NewReader(compressFile)
    left , offset  := readTable(*br , encodeMap)
    for idx , v  := range encodeMap {
        decodeMap[v] = idx
    }
    // 解码string暂存区
    var buff bytes.Buffer
    // 编码bytes暂存区
    codeBuff := make([]byte,contentBuffer)
    codeLen , _  := compressFile.ReadAt(codeBuff,int64(offset))
    //遍历解码 , 读取比特
    for i := 0 ; i < codeLen  ; i ++ {
        //对每个byte单独进行位运算转string
        perByte := codeBuff[i]
        for j := 0 ; j < 8 ; j ++ {
            //与运算
            buff.WriteString(strconv.Itoa( int( (perByte >> j) & 1 ) ) )
        }
    }
    // 对照码表,解码string , 对8取余目的是解决正好读满8个bit的情况发生
    contentStr  := buff.String()[:buff.Len() - (8 - left) % 8]
    bytes := make([]byte, 0 )
    //用切片读contenStr即可
    for star , end := 0 , 1 ; end <= len(contentStr) ; {
        charValue , ok := decodeMap[contentStr[star:end]]
        if ok {
            bytes = append(bytes, byte(charValue))
            star = end
        }
        end ++
    }

    depressFile , _ := os.Create(depressPath)
    depressFile.Write(bytes)
    depressFile.Close()
}

func readTable( br bufio.Reader , encodeMap map[int]string) (int ,int )  {
    lineStr ,  _  , _ := br.ReadLine()
    lines , _ := strconv.Atoi(string(lineStr))
    for i := 0 ; i < lines - 1  ; i ++ {
        lineContent , _ , _:= br.ReadLine()
        kvArr := strings.Split(string(lineContent),":")
        k , v := kvArr[0] , kvArr[1]
        kNum , _:= strconv.Atoi(k)
        encodeMap[kNum] = v
    }
    leftStr , _ ,_:= br.ReadLine()
    left , _ := strconv.Atoi(string(leftStr))
    return left , br.Size() - br.Buffered()
}

func MD5Bytes(s []byte) string {
    ret := md5.Sum(s)
    return hex.EncodeToString(ret[:])
}

func MD5File(file string) (string, error) {
    data, err := ioutil.ReadFile(file)
    if err != nil {
        return "", err
    }
    return MD5Bytes(data), nil
}

四、接下来按代码的实现顺序,分模块讲解代码的各个部分

  1. 统计词频,压缩的基石
    getFrequencyMap方法,使用map记录文件中字符对应ASCII码出现的次数,因为golang中没有char类型,这里用int替代
func getFrequencyMap( file *os.File ) map[int]int {
    imap := make(map[int]int)
    // 读入文件数据,readline 记入map中,统计频次
    reader := bufio.NewReader(file)
    buffer := make([]byte,contentBuffer)
    readCount , _ := reader.Read(buffer)
    for i := 0 ; i < readCount ; i ++ {
        imap[int(buffer[i])] ++
    }
    return imap
}

构造哈夫曼树的节点,顺便按照词频给这些节点进行排序

    plist := make(treeHeap,0)
    // 遍历map ,按频率排序
    for k , v := range imap {
        plist = append(plist, &TreeNode{Val: k,Times: v})
    }
    sort.Sort(plist)

其中heap堆的实现如下,实现了heap接口的结构体既是可排序数组,也可以当作优先队列来使用

type TreeNode struct {
    Val int
    Times int
    Left *TreeNode
    Right *TreeNode
}

type  treeHeap []*TreeNode

func (p treeHeap) Less( i , j int) bool {
    return p[i].Times <= p[j].Times
}

func (p treeHeap) Len() int {
    return len(p)
}

func (p treeHeap) Swap(i , j int) {
    p[i] , p[j] = p[j] , p[i]
}

func (p *treeHeap) Push(node interface{}) {
    *p = append(*p, node.(*TreeNode))
}

func (p *treeHeap) Pop()  interface{} {
    n := len(*p)
    t := (*p)[n-1]
    *p = (*p)[:n-1]
    return t
}
  1. 使用优先队列构造哈夫曼树
    使用递归也可以生成哈夫曼树,但是不能保证树的结构是最优的,我们需要保证构建的哈夫曼树具有最小权值路径和,这样才能使得压缩的效率最大化,从代码可以清晰的看到,实现了堆接口后的哈夫曼树构造非常简单。因为heap是优先队列的缘故,每次插入都会按Times(频次)升序排序保证下次合成的两两节点权值之和是所有节点中最小的,plist[0]即为构造完成哈夫曼树的根节点.
func initHuffmanTree(plist treeHeap) *TreeNode {
    // 使用优先队列构造最小路径权值哈夫曼树
    heap.Init(&plist)
    for plist.Len() > 1 {
        t1 := heap.Pop(&plist).(*TreeNode)
        t2 := heap.Pop(&plist).(*TreeNode)
        root := &TreeNode{Times: t1.Times + t2.Times}
        if t1.Times > t2.Times {
            root.Right  , root.Left = t1 , t2
        }else {
            root.Right , root.Left = t2 , t1
        }
        heap.Push(&plist,root)
    }
    return plist[0]
}
  1. 使用回溯法得到经过哈夫曼树编码以后的table
    encodeMap中key为字符的ASCII码值,用int表示,value为哈夫曼树从根节点到叶子节点的左右路径编码值,用string表示,其中tmp记录递归过程中所走过的路径,往左走用'0'表示,往右走用'1'表示,最后再转换成string就得到了对应ASCII码的编码值。
func createEncodingTable(node *TreeNode , encodeMap map[int]string )  {
    /*  思路:回溯遍历二叉树,用byte记录0 1🐎,遇到叶子节点就转换成string存入码表
        遍历顺序:根左右
    */
    tmp := make([]byte,0)
    var depth func (treeNode *TreeNode)
    depth = func (root *TreeNode) {
        //如果已经遍历到空,返回
        if root == nil {
            return
        }
        //如果遍历到的是叶子节点 , byte转换成string,入表
        if root.Left == nil && root.Right == nil {
            encodeMap[root.Val] = string(tmp)
        }
        // 如果是普通节点,左右递归回溯即可 #规则: 左0右1
        tmp = append(tmp ,'0' )
        depth(root.Left)
        tmp[len(tmp)-1] = '1'
        depth(root.Right)
        tmp = tmp[:len(tmp)-1]
    }
    depth(node)
}
  1. 我们来到了最关键的一步,这里将使用编码表将字符对应的哈夫曼编码转换成bit位存入文件中,转换过程使用了位运算,代码模块较长,先将代码贴出,再分块讲解
func encoding(inPath string, outPath string , encodeMap map[int]string) {
    /*1. 一次性读入文件内容*/
    inFile ,err  := os.Open(inPath)
    defer inFile.Close()
    if err != nil {
        return
    }
    reader := bufio.NewReader(inFile)
    fileContent := make([]byte,contentBuffer)
    count , _ := reader.Read(fileContent)
    var buff bytes.Buffer
    // string编码
    for i := 0 ; i < count ; i ++  {
        v := fileContent[i]
        if code , ok := encodeMap[int(v)] ; len(code)!= 0  && ok  {
            buff.WriteString(code)
        }
    }
    res := make([]byte,0)
    var buf byte = 0
    // bit编码
    // 记录bit剩余位,很简单只要对buff.bytes取长度对8取余即可
    for idx , bit := range buff.Bytes() {
        //每八个位使用一个byte读取,结果存入res数组即可
        pos := idx % 8
        if pos == 0 && idx > 0 {
            res = append(res, buf)
            buf = 0
        }
        if bit == '1' {
            buf  |=   1 <<  pos
        }
    }
    // 这个left是剩余待处理的位数
    left := buff.Len() % 8
    res = append(res, buf)
    // 先将码表和left数写入文件,解码时在开头读取
    writeTable(outPath,encodeMap,left)
    outFile , err := os.OpenFile(outPath,os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil{
        return
    }
    wcount , err := outFile.Write(res)
    if err != nil {
        fmt.Println(wcount)
        return
    }
}

首先是读入文件,将每一个字符都转换成被哈夫曼编码过的string,比如a转换成"01"

    /*1. 一次性读入文件内容*/
    inFile ,err  := os.Open(inPath)
    defer inFile.Close()
    if err != nil {
        return
    }
    reader := bufio.NewReader(inFile)
    fileContent := make([]byte,contentBuffer)
    count , _ := reader.Read(fileContent)
    var buff bytes.Buffer
    //string编码
    for i := 0 ; i < count ; i ++  {
        v := fileContent[i]
        if code , ok := encodeMap[int(v)] ; len(code)!= 0  && ok  {
            buff.WriteString(code)
        }
    }

其次,将文件中的每一个字符根据哈夫曼编码Map对应的string转换成一个个bit存入byte数组中

    res := make([]byte,0)
    var buf byte = 0
    // bit编码
    // 记录bit剩余位,很简单只要对buff.bytes取长度对8取余即可
    for idx , bit := range buff.Bytes() {
        //每八个位使用一个byte读取,结果存入res数组即可
        pos := idx % 8
        if pos == 0 && idx > 0 {
            res = append(res, buf)
            buf = 0
        }
        if bit == '1' {
            buf  |=   1 <<  pos
        }
    }
    // 这个left是剩余待处理的位数
    left := buff.Len() % 8
    res = append(res, buf)

这里详细讲解存入过程 :
每个byte的初始状态都是00000000(八位) , 假设要存入的字符串为"0111011110101"(13位),那么我们需要两个byte来记录字符串的信息,从前往后遍历字符串,用pos记录当前遍历到的字符串下标,读到'1'时,对byte的对应比特位进行或操作,这样进行八次读取,第一个byte就变成了011101111的倒序排列byte1 : 11101110,因为后续我们读入的时候也是倒序读入,所以这样做是可行的,还剩最后五位"10101"用一个新的byte来记录,同理byte2 : 00010101
byte2的前三位是空余位,读入的时候我们并不考虑,所以用left来记录最后一个byte要读入的bit位数即可。left = 13 % 8 = 5 。 我们将所有经过位运算的byte存入byte数组,先保留在内存中,等待完成文件头的处理,再把编码集写入文件,就完成了最重要的一步,哈夫曼编码压缩。

  1. 将table写入文件头,将编码结果写入文件
    writeTable函数,我们记录map的长度加上left位的长度,这些都是重要的信息,每个键值对用换行隔开,key和value用":"做分割,以备在读入的时候解析。
func writeTable(path string, codeMap map[int]string , left int) {
    file ,err := os.Create(path)
    if err != nil {
        return
    }
    // 第一行,写入文件头的长度
    var buff bytes.Buffer
    buff.WriteString(strconv.Itoa(len(codeMap)+1) + "\n")
    for k , v := range codeMap {
        buff.WriteString(strconv.Itoa(k) + ":" + v + "\n")
    }
    buff.WriteString(strconv.Itoa(left) + "\n")
    file.WriteString(buff.String())
    file.Close()
}

接下来,我们用追加写入模式打开文件,写入编码结果即可

 outFile , err := os.OpenFile(outPath,os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
    if err != nil{
        return
    }
    wcount , err := outFile.Write(res)
    if err != nil {
        fmt.Println(wcount)
        return
    }
  1. 读取文件,进行解压缩
    这一步就很简单了。有了文件头中的编码表table,以及我们之前做过的位运算基础,倒序操作一遍即可
func depress( inPath , depressPath string) {
    // originPath 原文件(或者可以传入码表), inPath 读入被压缩的文件 , depressPath 还原后的输出路径
    encodeMap := make(map[int]string)
    decodeMap := make(map[string]int)
    //2.读入压缩文件
    compressFile , _ := os.Open(inPath)
    // br 读取文件头 ,返回偏移量
    br := bufio.NewReader(compressFile)
    left , offset  := readTable(*br , encodeMap)
    for idx , v  := range encodeMap {
        decodeMap[v] = idx
    }
    // 解码string暂存区
    var buff bytes.Buffer
    // 编码bytes暂存区
    codeBuff := make([]byte,contentBuffer)
    codeLen , _  := compressFile.ReadAt(codeBuff,int64(offset))
    //遍历解码 , 读取比特
    for i := 0 ; i < codeLen  ; i ++ {
        //对每个byte单独进行位运算转string
        perByte := codeBuff[i]
        for j := 0 ; j < 8 ; j ++ {
            //与运算
            buff.WriteString(strconv.Itoa( int( (perByte >> j) & 1 ) ) )
        }
    }
    // 对照码表,解码string , 对8取余目的是解决正好读满8个bit的情况发生
    contentStr  := buff.String()[:buff.Len() - (8 - left) % 8]
    bytes := make([]byte, 0 )
    //用切片读contenStr即可
    for star , end := 0 , 1 ; end <= len(contentStr) ; {
        charValue , ok := decodeMap[contentStr[star:end]]
        if ok {
            bytes = append(bytes, byte(charValue))
            star = end
        }
        end ++
    }

    depressFile , _ := os.Create(depressPath)
    depressFile.Write(bytes)
    depressFile.Close()
}

其中的重点就是位运算操作,将bit信息还原成string,再到编码table中找到对应的ASCII码值

//遍历解码 , 读取比特
    for i := 0 ; i < codeLen  ; i ++ {
        //对每个byte单独进行位运算转string
        perByte := codeBuff[i]
        for j := 0 ; j < 8 ; j ++ {
            //与运算
            buff.WriteString(strconv.Itoa( int( (perByte >> j) & 1 ) ) )
        }
    }
// 对照码表,解码string , 对8取余目的是解决正好读满8个bit的情况发生
    contentStr  := buff.String()[:buff.Len() - (8 - left) % 8]
//用切片读contenStr即可
    for star , end := 0 , 1 ; end <= len(contentStr) ; {
        charValue , ok := decodeMap[contentStr[star:end]]
        if ok {
            bytes = append(bytes, byte(charValue))
            star = end
        }
        end ++
    }

详情参见注释,在此不做赘述

  1. 最后一个细节,在读取文件时,如何分别读入文件头和压缩内容
    我这里采用的是使用io.reader,按行号读取文件头后再计算已经读入的偏移量,最后使用file.readAt方法来接着读入文件内容
    readTable 调用点: 返回offset偏移量在最后ReadAt方法中使用
left , offset  := readTable(*br , encodeMap)
    for idx , v  := range encodeMap {
        decodeMap[v] = idx
    }
    // 解码string暂存区
    var buff bytes.Buffer
    // 编码bytes暂存区
    codeBuff := make([]byte,contentBuffer)
    codeLen , _  := compressFile.ReadAt(codeBuff,int64(offset))

readTable方法:

func readTable( br bufio.Reader , encodeMap map[int]string) (int ,int )  {
    lineStr ,  _  , _ := br.ReadLine()
    lines , _ := strconv.Atoi(string(lineStr))
    for i := 0 ; i < lines - 1  ; i ++ {
        lineContent , _ , _:= br.ReadLine()
        kvArr := strings.Split(string(lineContent),":")
        k , v := kvArr[0] , kvArr[1]
        kNum , _:= strconv.Atoi(k)
        encodeMap[kNum] = v
    }
    leftStr , _ ,_:= br.ReadLine()
    left , _ := strconv.Atoi(string(leftStr))
    return left , br.Size() - br.Buffered()
}

到此为止,实现哈夫曼编码压缩文件的所有方法和细节已经展示完毕,解决了我开头所述的三个痛点,即哈夫曼编码的代码实现,将编码转换成bit压缩至文件的实现,文件头读取的过程,如有不足之处烦请各位指教。

以上内容皆为本人原创,转载请注明出处