哈夫曼树

/**
* 哈夫曼树
* author：JetWu
* date：2020.02.09
 */

package huffmanTree

import (
	"errors"
	"fmt"
)

/**
* 哈夫曼树节点结构
**/
type HuffmanTreeNode struct {
	data   byte //数据
	weight int  //权重
	index  int  //当前节点索引
	parent int  //父节点索引
	left   int  //左子节点索引
	right  int  //右子节点索引
}

/**
* 哈夫曼树结构
**/
type HuffmanTree struct {
	slice []*HuffmanTreeNode //节点数组，最后一个元素即是根节点
	size  int                //叶子节点个数
}

/**
* 构造哈夫曼树
**/
func CreateHuffmanTree(data []byte, weight []int) *HuffmanTree {
	if len(data) < 2 || len(data) != len(weight) {
		return nil
	}
	count := len(data)
	hft := &HuffmanTree{
		slice: make([]*HuffmanTreeNode, count),
		size:  count,
	}
	mh := NewMinHeap()
	for i := 0; i < len(data); i++ { //将所有节点放入最小堆
		htNode := &HuffmanTreeNode{
			data:   data[i],
			weight: weight[i],
			index:  i,
			parent: -1,
			left:   -1,
			right:  -1,
		}
		hft.slice[i] = htNode
		mh.Insert(htNode)
	}
	for mh.Count() > 1 { //每次取出权重最小的两个节点，合并成树，将树根节点插入最小堆
		htNode1, _ := mh.RemoveMin()
		htNode2, _ := mh.RemoveMin()
		index := len(hft.slice)
		htNode := &HuffmanTreeNode{
			weight: htNode1.weight + htNode2.weight,
			index:  index,
			parent: -1,
			left:   htNode1.index,
			right:  htNode2.index,
		}
		hft.slice = append(hft.slice, htNode)
		htNode1.parent, htNode2.parent = index, index
		mh.Insert(htNode)
	}
	return hft
}

/**
* 打印所有字符及其对应编码
**/
func (hft *HuffmanTree) Print() {
	b := make([]byte, 1)
	for i := 0; i < hft.size; i++ {
		fmt.Print(string(hft.slice[i].data), ": ")
		b[0] = hft.slice[i].data
		code, _ := hft.Encode(string(b))
		fmt.Println(code)
	}
}

/**
* 打印所有节点信息
**/
func (hft *HuffmanTree) Print2() {
	for _, d := range hft.slice {
		fmt.Println("data:", string(d.data), ", weight:", d.weight, ", index:", d.index, ", parent:", d.parent, ", left:", d.left, ", right:", d.right)
	}
}

/**
* 编码
**/
func (hft *HuffmanTree) Encode(data string) (string, error) {
	if len(data) == 0 {
		return "", errors.New("数据为空")
	}
	if hft == nil {
		return "", errors.New("哈弗曼树为空")
	}
	bData := []byte(data)
	var code []byte
	as := NewArrayStack()
	for _, d := range bData {
		found := false
		for i := 0; i < hft.size; i++ {
			if hft.slice[i].data == d {
				found = true
				ptr := hft.slice[i]
				lastIndex := ptr.index
				for ptr.parent != -1 {
					ptr = hft.slice[ptr.parent]
					if ptr.left == lastIndex { //左0
						as.Push(‘0‘)
					} else { //右1
						as.Push(‘1‘)
					}
					lastIndex = ptr.index
				}
			}
		}
		if !found {
			return "", errors.New("编码不存在：" + string(d))
		}
		for !as.IsEmpty() { //先进后出，达到倒序效果
			b, _ := as.Pop()
			code = append(code, b)
		}
	}
	return string(code), nil
}

/**
* 解码
**/
func (hft *HuffmanTree) Decode(code string) (string, error) {
	if len(code) == 0 {
		return "", errors.New("编码为空")
	}
	if hft == nil {
		return "", errors.New("哈弗曼树为空")
	}
	bCode := []byte(code)
	var data []byte
	root := hft.slice[len(hft.slice)-1] //根节点
	ptr := root
	for _, c := range bCode {
		if c == ‘0‘ {
			ptr = hft.slice[ptr.left]
		} else if c == ‘1‘ {
			ptr = hft.slice[ptr.right]
		} else {
			return "", errors.New("编码错误：" + string(c))
		}
		if ptr.left == -1 && ptr.right == -1 { //叶子节点
			data = append(data, ptr.data)
			ptr = root
		}
	}
	if ptr != root {
		return "", errors.New("编码错误")
	}
	return string(data), nil
}

/** 以下为哈夫曼树构造过程使用的最小堆 **/

/**
* 最小堆结构
**/
type MinHeap struct {
	slice []*HuffmanTreeNode
	size  int //最小堆大小
}

/**
* 创建一个空最小堆
**/
func NewMinHeap() *MinHeap {
	return &MinHeap{
		slice: make([]*HuffmanTreeNode, 10),
		size:  0,
	}
}

/**
* 最小堆容量
**/
func (mh *MinHeap) Count() int {
	return mh.size
}

/**
* 向下调整（最小堆初始化、删除时调用）
**/
func (mh *MinHeap) shiftDown(start int) {
	i, j := start, start*2+1 //j为i的左子节点
	temp := mh.slice[i]
	for j < mh.size {
		if j < mh.size-1 && mh.slice[j+1].weight < mh.slice[j].weight { //比较获得左右子节点的最小值编号
			j++
		}
		if temp.weight <= mh.slice[j].weight { //左右子节点都比父节点大
			break
		} else { //继续向下比较
			mh.slice[i] = mh.slice[j]
			i, j = j, j*2+1
		}
	}
	mh.slice[i] = temp
}

/**
* 向上调整（最小堆插入时调用）
**/
func (mh *MinHeap) siftUp(start int) {
	if start > mh.size-1 {
		return
	}
	j, i := start, (start-1)/2 //i为j的父节点
	temp := mh.slice[j]
	for j > 0 {
		if temp.weight >= mh.slice[i].weight { //子节点大于父节点
			break
		} else {
			mh.slice[j] = mh.slice[i]
			j, i = i, (i-1)/2
		}
	}
	mh.slice[j] = temp
}

/**
* 插入
**/
func (mh *MinHeap) Insert(data *HuffmanTreeNode) {
	//将插入元素放置在最后节点
	mh.size++
	if mh.size <= len(mh.slice) {
		mh.slice[mh.size-1] = data
	} else {
		mh.slice = append(mh.slice, data)
	}
	//向上调整
	mh.siftUp(mh.size - 1)
}

/**
* 删除
**/
func (mh *MinHeap) RemoveMin() (*HuffmanTreeNode, error) {
	if mh.size < 1 {
		return nil, errors.New("最小堆为空")
	}
	min := mh.slice[0]
	mh.size--
	//使用最后一个节点替换第一个节点
	mh.slice[0] = mh.slice[mh.size]
	//向下调整
	mh.shiftDown(0)
	return min, nil
}

/** 以下为哈夫曼树编码过程使用的栈 **/

/**
* 数组栈结构
**/
type ArrayStack struct {
	slice []byte
	head  int //栈顶位置：slice最后一个有效元素的索引位置
}

/**
* 创建空栈
**/
func NewArrayStack() *ArrayStack {
	return &ArrayStack{
		slice: make([]byte, 10),
		head:  -1,
	}
}

/**
* 判断栈是否为空
**/
func (as *ArrayStack) IsEmpty() bool {
	return as.head == -1
}

/**
* 入栈
**/
func (as *ArrayStack) Push(data byte) bool {
	as.head++
	if as.head+1 <= len(as.slice) {
		as.slice[as.head] = data
	} else {
		as.slice = append(as.slice, data)
	}
	return true
}

/**
* 出栈
**/
func (as *ArrayStack) Pop() (byte, error) {
	if as.head == -1 {
		return 0, errors.New("栈为空")
	}
	data := as.slice[as.head]
	as.head--
	return data, nil
}