com.firefly.utils.codec.HuffmanCodec Maven / Gradle / Ivy
package com.firefly.utils.codec;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
public class HuffmanCodec implements Serializable{
private static final long serialVersionUID = -5318250039712365557L;
abstract public static class HuffmanTree implements Comparable>, Serializable {
private static final long serialVersionUID = -5354103251920897803L;
public final Long frequency; // the frequency of this tree
public HuffmanTree(Long frequency) {
this.frequency = frequency;
}
@Override
public int compareTo(HuffmanTree tree) {
return frequency.compareTo(tree.frequency);
}
}
public static class HuffmanLeaf extends HuffmanTree {
private static final long serialVersionUID = -8197406618091612264L;
public final T value;
public HuffmanLeaf(Long freq, T value) {
super(freq);
this.value = value;
}
@Override
public String toString() {
return "HuffmanLeaf [value=" + value + ", frequency=" + frequency
+ "]";
}
}
public static class HuffmanNode extends HuffmanTree {
private static final long serialVersionUID = -4581114135719242316L;
public final HuffmanTree left, right; // subtrees
public HuffmanNode(HuffmanTree left, HuffmanTree right) {
super(left.frequency + right.frequency);
this.left = left;
this.right = right;
}
}
public static class HuffmanCode implements Serializable {
private static final long serialVersionUID = -4696130695208200688L;
public final Long frequency;
public final int length;
public final String bits;
public final BitBuilder bitBuilder;
public HuffmanCode(Long frequency, String bits) {
super();
this.frequency = frequency;
this.length = bits.length();
this.bits = bits;
bitBuilder = new BitBuilder(length);
for (int i = 0; i < bits.length(); i++) {
bitBuilder.append(bits.charAt(i) == '1');
}
}
@Override
public String toString() {
return "HuffmanCode [frequency=" + frequency + ", length=" + length
+ ", bits=" + bits + ", getBytes()="
+ Arrays.toString(getBytes()) + "]";
}
public byte[] getBytes() {
return bitBuilder.toByteArray();
}
}
public static class BitBuilder extends BitSet{
private static final long serialVersionUID = 4678685861273345213L;
private int length;
private int index;
public BitBuilder() {
super();
}
public BitBuilder(int len) {
super(len);
}
public int getLength() {
return length;
}
public BitBuilder append(boolean value) {
set(index, value);
index++;
length++;
return this;
}
}
private Map codecMap;
private HuffmanTree huffmanTree;
public HuffmanCodec() {}
public HuffmanCodec(T[] elements) {
huffmanTree = buildHuffmanTree(elements);
codecMap = buildHuffmanCodeMap(huffmanTree);
}
public Map getCodecMap() {
return codecMap;
}
public HuffmanTree getHuffmanTree() {
return huffmanTree;
}
public void setCodecMap(Map codecMap) {
this.codecMap = codecMap;
}
public void setHuffmanTree(HuffmanTree huffmanTree) {
this.huffmanTree = huffmanTree;
}
public BitBuilder encode(T[] elements) {
BitBuilder bits = new BitBuilder();
for(T t : elements) {
HuffmanCode code = codecMap.get(t);
for (int j = 0; j < code.length; j++) {
bits.append(code.bitBuilder.get(j));
}
}
return bits;
}
public List decode(BitBuilder bits) {
List elements = new ArrayList<>();
HuffmanTree currentNode = huffmanTree;
for (int i = 0; i < bits.getLength(); i++) {
if(currentNode instanceof HuffmanNode) {
HuffmanNode node = (HuffmanNode)currentNode;
currentNode = bits.get(i) ? node.right : node.left;
}
if(currentNode instanceof HuffmanLeaf) {
HuffmanLeaf leaf = (HuffmanLeaf) currentNode;
elements.add(leaf.value);
currentNode = huffmanTree;
}
}
return elements;
}
public static Map buildHuffmanCodeMap(HuffmanTree tree) {
Map map = new HashMap<>();
StringBuilder bits = new StringBuilder();
_buildHuffmanCodeMap(tree, map, bits);
return map;
}
private static void _buildHuffmanCodeMap(HuffmanTree tree, Map map, StringBuilder bits) {
if(tree instanceof HuffmanLeaf) {
HuffmanLeaf leaf = (HuffmanLeaf)tree;
HuffmanCode code = new HuffmanCode(leaf.frequency, bits.toString());
map.put(leaf.value, code);
} else if(tree instanceof HuffmanNode) {
HuffmanNode node = (HuffmanNode) tree;
// traverse left
bits.append('0');
_buildHuffmanCodeMap(node.left, map, bits);
bits.deleteCharAt(bits.length()-1);
// traverse right
bits.append('1');
_buildHuffmanCodeMap(node.right, map, bits);
bits.deleteCharAt(bits.length()-1);
}
}
public static HuffmanTree buildHuffmanTree(T[] elements) {
Map frequencyMap = new HashMap();
for(T t : elements) {
Long num = frequencyMap.get(t);
if(num == null) {
frequencyMap.put(t, 1L);
} else {
frequencyMap.put(t, num + 1L);
}
}
// initially, we have a forest of leaves
PriorityQueue> trees = new PriorityQueue<>();
for(Map.Entry entry : frequencyMap.entrySet()) {
if(entry.getValue() != null && entry.getValue() > 0) {
trees.offer(new HuffmanLeaf(entry.getValue(), entry.getKey()));
}
}
// loop until there is only one tree left
while (trees.size() > 1) {
// two trees with least frequency
HuffmanTree a = trees.poll();
HuffmanTree b = trees.poll();
// put into new node and re-insert into queue
trees.offer(new HuffmanNode(a, b));
}
return trees.poll();
}
}