Information Theory - HW1 (tmp)

import sys
import heapq as hp

''' class for heap tree '''


class Node:
    def __init__(self, index, freq):
        self.index = index
        self.freq = freq
        self.left = None
        self.right = None

    def __lt__(self, other):
        return self.freq < other.freq


def rec(elem, len, res):
    if elem.index != None:
        res.append((elem.index, len))
        return
    rec(elem.left, len+1, res)
    rec(elem.right, len+1, res)


def huffman(arr):
    # create array of tuples (frequency, index)
    # [(0.009047, 6), (0.021247, 3), (0.041388, 7), (0.045914, 4), (0.047731, 5)...]
    arr = [(arr[x], x) for x in range(len(arr))]
    # print(sorted(arr, key=lambda x: x[0]))

    # create queue
    pq = []
    for key in arr:
        node = Node(freq=key[0], index=key[1])
        hp.heappush(pq, node)

    # create tree
    while len(pq) > 1:
        f = hp.heappop(pq)
        s = hp.heappop(pq)

        merged = Node(freq=f.freq + s.freq, index=None)
        merged.left = f
        merged.right = s

        hp.heappush(pq, merged)

    # evaluate code lengths
    res = []
    rec(hp.heappop(pq), 0, res)
    # print(sorted(res, key=lambda x: x[1]))
    return res


def craft_value(arr):
    # arr = [(freq, len)]
    sum = 0
    for _, len in arr:
        sum += 2**(0-len)
    return sum


def build_code(arr):
    # Create tuple of (length, index, code)
    for i in range(n):
        arr[i] = (arr[i][1], arr[i][0], '')

    # sort by length of codes
    arr = sorted(arr, key=lambda x: x[0])

    dlm = ''
    for i in range(len(arr)):
        tmp = list(arr[i])

        tmp[2] = dlm + '0' * (arr[i][0] - len(dlm))
        dlm = "{0:b}".format((int(tmp[2], 2) + 1))
        dlm = '0' * (len(tmp[2]) - len(dlm)) + dlm

        arr[i] = tmp

    # sort by length of codes
    arr = sorted(arr, key=lambda x: x[1])

    return arr


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage : python3 Huffman.py src-fname dest-fname")
        exit()

    src_fname = sys.argv[1]
    dst_fname = sys.argv[2]

    with open(src_fname, 'r') as f:
        n = int(f.readline())
        arr = list(map(float, f.readline().split()))

    # find length of codes with probabilities
    arr = huffman(arr)

    # calculate craft value
    cv = craft_value(arr)
    if cv > 1:
        exit()

    # build code
    arr = build_code(arr)

    # Write Result
    with open(dst_fname, 'w+') as df:
        # Craft Value
        df.write('{}\n'.format(cv))

        for len, ind, code in arr:
            df.write('{}\n'.format(code))