aboutsummaryrefslogtreecommitdiff
path: root/src/huffmantree.nim
diff options
context:
space:
mode:
Diffstat (limited to 'src/huffmantree.nim')
-rw-r--r--src/huffmantree.nim31
1 files changed, 4 insertions, 27 deletions
diff --git a/src/huffmantree.nim b/src/huffmantree.nim
index 44c9990..0266dfb 100644
--- a/src/huffmantree.nim
+++ b/src/huffmantree.nim
@@ -30,35 +30,21 @@ type HuffmanTreeNode*[T: SomeUnsignedInt] = ref object
30 maxChildValue: T 30 maxChildValue: T
31 of leaf: 31 of leaf:
32 value*: T 32 value*: T
33 weight: int
34 33
35proc huffmanBranch*[T](left, right: HuffmanTreeNode[T]): HuffmanTreeNode[T] = 34proc huffmanBranch*[T](left, right: HuffmanTreeNode[T]): HuffmanTreeNode[T] =
36 HuffmanTreeNode[T]( 35 HuffmanTreeNode[T](
37 kind: branch, left: left, right: right, 36 kind: branch, left: left, right: right,
38 maxChildValue: max(left.maxValue(), right.maxValue()), 37 maxChildValue: max(left.maxValue(), right.maxValue()))
39 weight: left.weight + right.weight)
40 38
41proc huffmanLeaf*[T](value: T, weight = 0): HuffmanTreeNode[T] = 39proc huffmanLeaf*[T](value: T): HuffmanTreeNode[T] =
42 HuffmanTreeNode[T](kind: leaf, value: value, weight: weight) 40 HuffmanTreeNode[T](kind: leaf, value: value)
43 41
44proc `==`*[T](a, b: HuffmanTreeNode[T]): bool = 42proc `==`*[T](a, b: HuffmanTreeNode[T]): bool =
45 if a.kind != b.kind or a.weight != b.weight: return false
46 case a.kind:
47 of branch: a.left == b.left and a.right == b.right
48 of leaf: a.value == b.value
49
50proc `~=`*[T](a, b: HuffmanTreeNode[T]): bool =
51 if a.kind != b.kind: return false 43 if a.kind != b.kind: return false
52 case a.kind: 44 case a.kind:
53 of branch: a.left ~= b.left and a.right ~= b.right 45 of branch: a.left == b.left and a.right == b.right
54 of leaf: a.value == b.value 46 of leaf: a.value == b.value
55 47
56proc `!~`*[T](a, b: HuffmanTreeNode[T]): bool =
57 not (a ~= b)
58
59proc `<`*[T](left, right: HuffmanTreeNode[T]): bool =
60 left.weight < right.weight
61
62proc maxValue*[T](node: HuffmanTreeNode[T]): T = 48proc maxValue*[T](node: HuffmanTreeNode[T]): T =
63 case node.kind: 49 case node.kind:
64 of branch: node.maxChildValue 50 of branch: node.maxChildValue
@@ -86,12 +72,3 @@ proc serialise*[T](tree: HuffmanTreeNode[T], bitWriter: BitWriter) =
86 bitWriter.writeBits(valueBitLength, node.value) 72 bitWriter.writeBits(valueBitLength, node.value)
87 bitWriter.writeBits(valueLengthFieldBitLength, valueBitLength.uint8) 73 bitWriter.writeBits(valueLengthFieldBitLength, valueBitLength.uint8)
88 writeNode(tree) 74 writeNode(tree)
89
90proc symbolQueue*[T](stats: CountTableRef[T]): HeapQueue[HuffmanTreeNode[T]] =
91 result = newHeapQueue[HuffmanTreeNode[T]]()
92 for item, count in stats.pairs: result.push(huffmanLeaf(item, count))
93
94proc buildHuffmanTree*[T: SomeUnsignedInt](stats: CountTableRef[T]): HuffmanTreeNode[T] =
95 var symbolQueue = symbolQueue(stats)
96 while symbolQueue.len > 1: symbolQueue.push(huffmanBranch(symbolQueue.pop(), symbolQueue.pop()))
97 symbolQueue[0]