* Improve efficiency of huffman coding

This commit is contained in:
Matthew Honnibal 2015-07-12 01:31:37 +02:00
parent ff9ff6f3fa
commit aad96fdb5c
1 changed files with 42 additions and 16 deletions

View File

@ -5,6 +5,8 @@ from libc.stdint cimport uint64_t
import numpy
cimport cython
cdef struct Node:
float prob
@ -12,6 +14,14 @@ cdef struct Node:
int right
cdef struct Code:
uint64_t bits
int length
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
cpdef list huffman_encode(float[:] probs):
assert len(probs) >= 3
@ -22,29 +32,34 @@ cpdef list huffman_encode(float[:] probs):
cdef int i = size - 1
cdef int j = 0
while i >= 0 or (j+1) < len(nodes):
while i >= 0 or (j+1) < nodes.size():
if i < 0:
cover_two_nodes(nodes, j)
j += 2
elif j >= len(nodes):
elif j >= nodes.size():
cover_two_words(nodes, i, i-1, probs[i]+probs[i-1])
i -= 2
elif i >= 1 and (j == len(nodes) or probs[i-1] < nodes[j].prob):
elif i >= 1 and (j == nodes.size() or probs[i-1] < nodes[j].prob):
cover_two_words(nodes, i, i-1, probs[i] + probs[i-1])
i -= 2
elif (j+1) < len(nodes) and nodes[j+1].prob < probs[i]:
elif (j+1) < nodes.size() and nodes[j+1].prob < probs[i]:
cover_two_nodes(nodes, j)
j += 2
else:
cover_one_word_one_node(nodes, j, i, probs[i])
i -= 1
j += 1
output = ['' for _ in range(len(probs))]
assign_codes(nodes, output, len(nodes) - 1, b'')
cdef vector[Code] codes
codes.resize(len(probs))
assign_codes(nodes, codes, len(nodes) - 1, b'')
output = []
for i in range(len(codes)):
out_str = '{0:b}'.format(codes[i].bits).rjust(codes[i].length, '0')
output.append(out_str)
return output
cdef int assign_codes(vector[Node]& nodes, list codes, int i, bytes path) except -1:
cdef int assign_codes(vector[Node]& nodes, vector[Code]& codes, int i, bytes path) except -1:
left_path = path + b'0'
right_path = path + b'1'
# Assign down left branch
@ -52,13 +67,17 @@ cdef int assign_codes(vector[Node]& nodes, list codes, int i, bytes path) except
assign_codes(nodes, codes, nodes[i].left, left_path)
else:
# Leaf on left
codes[-(nodes[i].left + 1)] = left_path
id_ = -(nodes[i].left + 1)
codes[id_].length = len(left_path)
codes[id_].bits = <uint64_t>int(left_path, 2)
# Assign down right branch
if nodes[i].right >= 0:
assign_codes(nodes, codes, nodes[i].right, right_path)
else:
# Leaf on right
codes[-(nodes[i].right + 1)] = right_path
id_ = -(nodes[i].right + 1)
codes[id_].length = len(right_path)
codes[id_].bits = <uint64_t>int(right_path, 2)
cdef int cover_two_nodes(vector[Node]& nodes, int j) nogil:
@ -69,19 +88,26 @@ cdef int cover_two_nodes(vector[Node]& nodes, int j) nogil:
nodes.push_back(node)
cdef int cover_one_word_one_node(vector[Node]& nodes, int j, int id_, float prob) except -1:
cdef int cover_one_word_one_node(vector[Node]& nodes, int j, int id_, float prob) nogil:
cdef Node node
# Encode leaves as negative integers, where the integer is the index of the
# word in the vocabulary.
leaf_id = - <int64_t>(id_ + 1)
new_prob = prob + nodes[j].prob
cdef int64_t leaf_id = - <int64_t>(id_ + 1)
cdef float new_prob = prob + nodes[j].prob
if prob < nodes[j].prob:
node = Node(left=leaf_id, right=j, prob=new_prob)
node.left = leaf_id
node.right = j
node.prob = new_prob
else:
node = Node(left=j, right=leaf_id, prob=new_prob)
node.left = j
node.right = leaf_id
node.prob = new_prob
nodes.push_back(node)
cdef int cover_two_words(vector[Node]& nodes, int id1, int id2, float prob) except -1:
cdef Node node = Node(left=-(id1 + 1), right=-(id2 + 1), prob=prob)
cdef int cover_two_words(vector[Node]& nodes, int id1, int id2, float prob) nogil:
cdef Node node
node.left = -(id1+1)
node.right = -(id2+1)
node.prob = prob
nodes.push_back(node)