From 9beec742dab151960aee8d2e0248f0b144334048 Mon Sep 17 00:00:00 2001
From: Vojtech Moravec <vojtech.moravec.st@vsb.cz>
Date: Tue, 24 Nov 2020 14:58:21 +0100
Subject: [PATCH] Add very compact method for saving Huffman trees.

---
 .../java/cz/it4i/qcmp/huffman/Huffman.java    | 40 ++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/main/java/cz/it4i/qcmp/huffman/Huffman.java b/src/main/java/cz/it4i/qcmp/huffman/Huffman.java
index 624fbf9..4601e97 100644
--- a/src/main/java/cz/it4i/qcmp/huffman/Huffman.java
+++ b/src/main/java/cz/it4i/qcmp/huffman/Huffman.java
@@ -1,5 +1,9 @@
 package cz.it4i.qcmp.huffman;
 
+import cz.it4i.qcmp.io.InBitStream;
+import cz.it4i.qcmp.io.OutBitStream;
+
+import java.io.IOException;
 import java.util.*;
 
 public class Huffman {
@@ -33,7 +37,7 @@ public class Huffman {
             parentB.setBit(0);
 
             final double mergedProbabilities = parentA.getProbability() + parentB.getProbability();
-            final HuffmanNode mergedNode = new HuffmanNode(mergedProbabilities, parentA, parentB);
+            final HuffmanNode mergedNode = HuffmanNode.constructWithProbability(parentA, parentB, mergedProbabilities);
             queue.add(mergedNode);
         }
         root = queue.poll();
@@ -124,4 +128,38 @@ public class Huffman {
         }
         return sortedHashMap;
     }
+
+    private void encodeHuffmanNode(final HuffmanNode node, final OutBitStream bitStream) throws IOException {
+        if (node.isLeaf()) {
+            bitStream.writeBit(1);
+            bitStream.write(node.getSymbol());
+        } else {
+            bitStream.writeBit(0);
+            encodeHuffmanNode(node.getSubNodeA(), bitStream);
+            encodeHuffmanNode(node.getSubNodeB(), bitStream);
+        }
+    }
+
+    private static HuffmanNode decodeHuffmanNode(final InBitStream bitStream) throws IOException {
+        if (bitStream.readBit()) // Leaf
+        {
+            return HuffmanNode.constructWithSymbol(null, null, bitStream.readValue());
+        } else {
+            final HuffmanNode nodeA = decodeHuffmanNode(bitStream);
+            nodeA.setBit(1);
+            final HuffmanNode nodeB = decodeHuffmanNode(bitStream);
+            nodeB.setBit(0);
+            return HuffmanNode.constructWithSymbol(nodeA, nodeB, -1);
+        }
+    }
+
+
+    public void saveHuffmanTree(final OutBitStream bitStream) throws IOException {
+        assert root != null : "The tree is not build.";
+        encodeHuffmanNode(root, bitStream);
+    }
+
+    public static HuffmanNode readHuffmanTree(final InBitStream bitStream) throws IOException {
+        return decodeHuffmanNode(bitStream);
+    }
 }
-- 
GitLab