Huffman Encryption

huffman_encryption.py

import graphviz
import heapq
import copy
import math

source_information = []
r = 2


def input_information():
    global source_information
    global r
    print("Enter the expansion source")
    r = float(input())
    print("Enter the probability of occurrence. (A/B, 0/* or */0 is end) ")
    cnt = 0

    while True:
        A, B = input().split("/")
        A = float(A)
        B = float(B)
        if A == 0 or B == 0:
            break
        cnt += 1
        source_information.append((A / B, [cnt]))

def generate_node_name(number):
    label_str = str(number[0])
    for now in number[1:]:
        label_str += ("," + str(now))
    return label_str

def generate_graph():
    global source_information
    global r
    graph = graphviz.Graph(format="png")
    data = copy.deepcopy(source_information)
    heapq.heapify(data)
    while len(data) > 1:
        cnt = 0
        now_data = []
        number = []
        sum_value = 0
        while cnt < r and len(data) > 0:
            now_data.append(heapq.heappop(data))
            cnt += 1
        for now in now_data:
            sum_value += now[0]
            number.extend(now[1])
        number.sort()
        next_node_name = generate_node_name(number)
        for idx, now in enumerate(now_data):
            graph.edge(generate_node_name(now[1]), next_node_name, label=str(idx))
        heapq.heappush(data, (sum_value, number))
    graph.view()

def about_encoding():
    input_information()
    generate_graph()

def calc_efficiency():
    global source_information
    global r
    print("Enter the expansion source")
    r = float(input())
    print("Enter the probability of occurrence. and length(A/B, 0/* or */0 is end) ")

    while True:
        A, B = input().split("/")
        A = float(A)
        B = float(B)
        encoding = int(input())
        if A == 0 or B == 0:
            break
        source_information.append((A, B, encoding))
    H = 0
    print("H &= ", end="")
    # tex output
    for now in source_information:
        print(r"- \dfrac{%d}{%d}\log_2 \dfrac{%d}{%d}" % (now[0], now[1], now[0], now[1]))
    print("&= ", end="")
    for now in source_information:
        temp = - (now[0] / now[1]) * math.log2(now[0] / now[1])
        print(" + %.5f" % temp)
        H += temp
    print("&= %.5f" % H, end="\n")
    L = 0
    print("L &= ", end="")
    for now in source_information:
        print(r"+ \dfrac{%d}{%d} \times %d" % (now[0], now[1], now[2]))
        L += (now[0] / now[1]) * now[2]
    print(r"&= %.5f" % L, end="\n")
    print(r"ans = \dfrac{H / L}{\log_2 r} = \dfrac{%.5f}{%.5f} = %f" % (H, L, (H / L) / math.log2(r)))


if __name__ == "__main__":
    select = int(input())
    if select == 0:
        input_information()
        generate_graph()
    else:
        calc_efficiency()