2015-07-11 16:23:35 +03:00
|
|
|
from libcpp.vector cimport vector
|
|
|
|
from libc.stdint cimport uint32_t
|
|
|
|
from libc.stdint cimport int64_t
|
|
|
|
from libc.stdint cimport uint64_t
|
|
|
|
|
|
|
|
import numpy
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
cimport cython
|
|
|
|
|
2015-07-11 16:23:35 +03:00
|
|
|
|
|
|
|
cdef struct Node:
|
|
|
|
float prob
|
|
|
|
int left
|
|
|
|
int right
|
|
|
|
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
cdef struct Code:
|
|
|
|
uint64_t bits
|
|
|
|
int length
|
|
|
|
|
|
|
|
|
|
|
|
@cython.boundscheck(False)
|
|
|
|
@cython.wraparound(False)
|
|
|
|
@cython.nonecheck(False)
|
2015-07-11 21:01:10 +03:00
|
|
|
cpdef list huffman_encode(float[:] probs):
|
2015-07-11 16:23:35 +03:00
|
|
|
assert len(probs) >= 3
|
|
|
|
|
|
|
|
output = numpy.zeros(shape=(len(probs),), dtype=numpy.uint64)
|
|
|
|
|
|
|
|
cdef int size = len(probs)
|
|
|
|
cdef vector[Node] nodes
|
|
|
|
cdef int i = size - 1
|
|
|
|
cdef int j = 0
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
while i >= 0 or (j+1) < nodes.size():
|
2015-07-11 16:23:35 +03:00
|
|
|
if i < 0:
|
2015-07-11 23:27:43 +03:00
|
|
|
cover_two_nodes(nodes, j)
|
2015-07-11 16:23:35 +03:00
|
|
|
j += 2
|
2015-07-12 02:31:37 +03:00
|
|
|
elif j >= nodes.size():
|
2015-07-11 23:27:43 +03:00
|
|
|
cover_two_words(nodes, i, i-1, probs[i]+probs[i-1])
|
2015-07-11 21:01:10 +03:00
|
|
|
i -= 2
|
2015-07-12 02:31:37 +03:00
|
|
|
elif i >= 1 and (j == nodes.size() or probs[i-1] < nodes[j].prob):
|
2015-07-11 23:27:43 +03:00
|
|
|
cover_two_words(nodes, i, i-1, probs[i] + probs[i-1])
|
2015-07-11 16:23:35 +03:00
|
|
|
i -= 2
|
2015-07-12 02:31:37 +03:00
|
|
|
elif (j+1) < nodes.size() and nodes[j+1].prob < probs[i]:
|
2015-07-11 23:27:43 +03:00
|
|
|
cover_two_nodes(nodes, j)
|
2015-07-11 16:23:35 +03:00
|
|
|
j += 2
|
|
|
|
else:
|
2015-07-11 23:27:43 +03:00
|
|
|
cover_one_word_one_node(nodes, j, i, probs[i])
|
2015-07-11 16:23:35 +03:00
|
|
|
i -= 1
|
|
|
|
j += 1
|
2015-07-12 02:31:37 +03:00
|
|
|
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)
|
2015-07-11 16:23:35 +03:00
|
|
|
return output
|
|
|
|
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
cdef int assign_codes(vector[Node]& nodes, vector[Code]& codes, int i, bytes path) except -1:
|
2015-07-11 21:01:10 +03:00
|
|
|
left_path = path + b'0'
|
|
|
|
right_path = path + b'1'
|
|
|
|
# Assign down left branch
|
2015-07-11 11:57:30 +03:00
|
|
|
if nodes[i].left >= 0:
|
2015-07-11 21:01:10 +03:00
|
|
|
assign_codes(nodes, codes, nodes[i].left, left_path)
|
2015-07-11 16:23:35 +03:00
|
|
|
else:
|
2015-07-11 21:01:10 +03:00
|
|
|
# Leaf on left
|
2015-07-12 02:31:37 +03:00
|
|
|
id_ = -(nodes[i].left + 1)
|
|
|
|
codes[id_].length = len(left_path)
|
|
|
|
codes[id_].bits = <uint64_t>int(left_path, 2)
|
2015-07-11 21:01:10 +03:00
|
|
|
# Assign down right branch
|
2015-07-11 11:57:30 +03:00
|
|
|
if nodes[i].right >= 0:
|
2015-07-11 21:01:10 +03:00
|
|
|
assign_codes(nodes, codes, nodes[i].right, right_path)
|
2015-07-11 16:23:35 +03:00
|
|
|
else:
|
2015-07-11 21:01:10 +03:00
|
|
|
# Leaf on right
|
2015-07-12 02:31:37 +03:00
|
|
|
id_ = -(nodes[i].right + 1)
|
|
|
|
codes[id_].length = len(right_path)
|
|
|
|
codes[id_].bits = <uint64_t>int(right_path, 2)
|
2015-07-11 11:57:30 +03:00
|
|
|
|
|
|
|
|
2015-07-11 23:27:43 +03:00
|
|
|
cdef int cover_two_nodes(vector[Node]& nodes, int j) nogil:
|
2015-07-11 16:23:35 +03:00
|
|
|
cdef Node node
|
|
|
|
node.left = j
|
|
|
|
node.right = j+1
|
|
|
|
node.prob = nodes[j].prob + nodes[j+1].prob
|
|
|
|
nodes.push_back(node)
|
2015-07-11 11:57:30 +03:00
|
|
|
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
cdef int cover_one_word_one_node(vector[Node]& nodes, int j, int id_, float prob) nogil:
|
2015-07-11 16:23:35 +03:00
|
|
|
cdef Node node
|
2015-07-11 11:57:30 +03:00
|
|
|
# Encode leaves as negative integers, where the integer is the index of the
|
|
|
|
# word in the vocabulary.
|
2015-07-12 02:31:37 +03:00
|
|
|
cdef int64_t leaf_id = - <int64_t>(id_ + 1)
|
|
|
|
cdef float new_prob = prob + nodes[j].prob
|
2015-07-11 11:57:30 +03:00
|
|
|
if prob < nodes[j].prob:
|
2015-07-12 02:31:37 +03:00
|
|
|
node.left = leaf_id
|
|
|
|
node.right = j
|
|
|
|
node.prob = new_prob
|
2015-07-11 11:57:30 +03:00
|
|
|
else:
|
2015-07-12 02:31:37 +03:00
|
|
|
node.left = j
|
|
|
|
node.right = leaf_id
|
|
|
|
node.prob = new_prob
|
2015-07-11 23:27:43 +03:00
|
|
|
nodes.push_back(node)
|
2015-07-11 11:57:30 +03:00
|
|
|
|
|
|
|
|
2015-07-12 02:31:37 +03:00
|
|
|
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
|
2015-07-11 16:23:35 +03:00
|
|
|
nodes.push_back(node)
|