com.firefly.utils.codec.HuffmanCodec Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of firefly-common Show documentation
There is a newer version: 5.0.2
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();
	}
}