Skip to content
Navigation Menu
{{ message }}
forked from keon/algorithms
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathhuffman_coding.py
More file actions
328 lines (265 loc) · 9.96 KB
/
Copy pathhuffman_coding.py
File metadata and controls
328 lines (265 loc) · 9.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
"""
Huffman coding is an efficient method of compressing data without losing information.
This algorithm analyzes the symbols that appear in a message.
Symbols that appear more often will be encoded as a shorter-bit string
while symbols that aren't used as much will be encoded as longer strings.
"""
from collections import defaultdict, deque
import heapq
class Node:
def __init__(self, frequency=0, sign=None, left=None, right=None):
self.frequency = frequency
self.sign = sign
self.left = left
self.right = right
def __lt__(self, other):
return self.frequency < other.frequency
def __gt__(self, other):
return self.frequency > other.frequency
def __eq__(self, other):
return self.frequency == other.frequency
def __str__(self):
return "<ch: {0}: {1}>".format(self.sign, self.frequency)
def __repr__(self):
return "<ch: {0}: {1}>".format(self.sign, self.frequency)
class HuffmanReader:
def __init__(self, file):
self.file = file
self.buffer = []
self.is_last_byte = False
def get_number_of_additional_bits_in_the_last_byte(self) -> int:
bin_num = self.get_bit() + self.get_bit() + self.get_bit()
return int(bin_num, 2)
def load_tree(self) -> Node:
"""
Load tree from file
:return:
"""
node_stack = deque()
queue_leaves = deque()
root = Node()
current_node = root
is_end_of_tree = False
while not is_end_of_tree:
current_bit = self.get_bit()
if current_bit == "0":
current_node.left = Node()
current_node.right = Node()
node_stack.append(current_node.right) # going to left node, right push on stack
current_node = current_node.left
else:
queue_leaves.append(current_node)
if node_stack:
current_node = node_stack.pop()
else:
is_end_of_tree = True
self._fill_tree(queue_leaves)
return root
def _fill_tree(self, leaves_queue):
"""
Load values to tree after reading tree
:param leaves_queue:
:return:
"""
leaves_queue.reverse()
while leaves_queue:
node = leaves_queue.pop()
s = int(self.get_byte(), 2)
node.sign = s
def _load_byte(self, buff_limit=8) -> bool:
"""
Load next byte is buffer is less than buff_limit
:param buff_limit:
:return: True if there is enough bits in buffer to read
"""
if len(self.buffer) <= buff_limit:
byte = self.file.read(1)
if not byte:
return False
i = int.from_bytes(byte, "big")
self.buffer.extend(list("{0:08b}".format(i)))
return True
def get_bit(self, buff_limit=8):
if self._load_byte(buff_limit):
bit = self.buffer.pop(0)
return bit
else:
return -1
def get_byte(self):
if self._load_byte():
byte_list = self.buffer[:8]
self.buffer = self.buffer[8:]
return "".join(byte_list)
else:
return -1
class HuffmanWriter:
def __init__(self, file):
self.file = file
self.buffer = ""
self.saved_bits = 0
def write_char(self, char):
self.write_int(ord(char))
def write_int(self, num):
bin_int = "{0:08b}".format(num)
self.write_bits(bin_int)
def write_bits(self, bits):
self.saved_bits += len(bits)
self.buffer += bits
while len(self.buffer) >= 8:
i = int(self.buffer[:8], 2)
self.file.write(bytes([i]))
self.buffer = self.buffer[8:]
def save_tree(self, tree):
"""
Generate and save tree code to file
:param tree:
:return:
"""
signs = []
tree_code = ""
def get_code_tree(T):
nonlocal tree_code
if T.sign is not None:
signs.append(T.sign)
if T.left:
tree_code += "0"
get_code_tree(T.left)
if T.right:
tree_code += "1"
get_code_tree(T.right)
get_code_tree(tree)
self.write_bits(tree_code + "1") # "1" indicates that tree ended (it will be needed to load the tree)
for int_sign in signs:
self.write_int(int_sign)
def _save_information_about_additional_bits(self, additional_bits: int):
"""
Overwrite first three bits in the file
:param additional_bits: number of bits that were appended to fill last byte
:return:
"""
self.file.seek(0)
first_byte_raw = self.file.read(1)
self.file.seek(0)
first_byte = "{0:08b}".format(int.from_bytes(first_byte_raw, "big"))
# overwrite first three bits
first_byte = first_byte[3:]
first_byte = "{0:03b}".format(additional_bits) + first_byte
self.write_bits(first_byte)
def close(self):
additional_bits = 8 - len(self.buffer)
if additional_bits != 8: # buffer is empty, no need to append extra "0"
self.write_bits("0" * additional_bits)
self._save_information_about_additional_bits(additional_bits)
class TreeFinder:
"""
Class to help find signs in tree
"""
def __init__(self, tree):
self.root = tree
self.current_node = tree
self.found = None
def find(self, bit):
"""
Find sign in tree
:param bit:
:return: True if sign is found
"""
if bit == "0":
self.current_node = self.current_node.left
elif bit == "1":
self.current_node = self.current_node.right
else:
self._reset()
return True
if self.current_node.sign is not None:
self._reset(self.current_node.sign)
return True
else:
return False
def _reset(self, found=""):
self.found = found
self.current_node = self.root
class HuffmanCoding:
def __init__(self):
pass
@staticmethod
def decode_file(file_in_name, file_out_name):
with open(file_in_name, "rb") as file_in, open(file_out_name, "wb") as file_out:
reader = HuffmanReader(file_in)
additional_bits = reader.get_number_of_additional_bits_in_the_last_byte()
tree = reader.load_tree()
HuffmanCoding._decode_and_write_signs_to_file(file_out, reader, tree, additional_bits)
print("File decoded.")
@staticmethod
def _decode_and_write_signs_to_file(file, reader: HuffmanReader, tree: Node, additional_bits: int):
tree_finder = TreeFinder(tree)
is_end_of_file = False
while not is_end_of_file:
bit = reader.get_bit()
if bit != -1:
while not tree_finder.find(bit): # read whole code
bit = reader.get_bit(0)
file.write(bytes([tree_finder.found]))
else: # There is last byte in buffer to parse
is_end_of_file = True
last_byte = reader.buffer
last_byte = last_byte[:-additional_bits] # remove additional "0" used to fill byte
for bit in last_byte:
if tree_finder.find(bit):
file.write(bytes([tree_finder.found]))
@staticmethod
def encode_file(file_in_name, file_out_name):
with open(file_in_name, "rb") as file_in, open(file_out_name, mode="wb+") as file_out:
signs_frequency = HuffmanCoding._get_char_frequency(file_in)
file_in.seek(0)
tree = HuffmanCoding._create_tree(signs_frequency)
codes = HuffmanCoding._generate_codes(tree)
writer = HuffmanWriter(file_out)
writer.write_bits("000") # leave space to save how many bits will be appended to fill the last byte
writer.save_tree(tree)
HuffmanCoding._encode_and_write_signs_to_file(file_in, writer, codes)
writer.close()
print("File encoded.")
@staticmethod
def _encode_and_write_signs_to_file(file, writer: HuffmanWriter, codes: dict):
sign = file.read(1)
while sign:
int_char = int.from_bytes(sign, "big")
writer.write_bits(codes[int_char])
sign = file.read(1)
@staticmethod
def _get_char_frequency(file) -> dict:
is_end_of_file = False
signs_frequency = defaultdict(lambda: 0)
while not is_end_of_file:
prev_pos = file.tell()
sign = file.read(1)
curr_pos = file.tell()
if prev_pos == curr_pos:
is_end_of_file = True
else:
signs_frequency[int.from_bytes(sign, "big")] += 1
return signs_frequency
@staticmethod
def _generate_codes(tree: Node) -> dict:
codes = dict()
HuffmanCoding._go_through_tree_and_create_codes(tree, "", codes)
return codes
@staticmethod
def _create_tree(signs_frequency: dict) -> Node:
nodes = [Node(frequency=frequency, sign=char_int) for char_int, frequency in signs_frequency.items()]
heapq.heapify(nodes)
while len(nodes) > 1:
left = heapq.heappop(nodes)
right = heapq.heappop(nodes)
new_node = Node(frequency=left.frequency + right.frequency, left=left, right=right)
heapq.heappush(nodes, new_node)
return nodes[0] # root
@staticmethod
def _go_through_tree_and_create_codes(tree: Node, code: str, dict_codes: dict):
if tree.sign is not None:
dict_codes[tree.sign] = code
if tree.left:
HuffmanCoding._go_through_tree_and_create_codes(tree.left, code + "0", dict_codes)
if tree.right:
HuffmanCoding._go_through_tree_and_create_codes(tree.right, code + "1", dict_codes)
You can’t perform that action at this time.