ACSL 2018 Compressed Trees Problem

ACSL2018CompressedTrees.py

#!/usr/bin/env python
import fileinput
from collections import Counter


def compress_list(_list, tree=None):
    """Compress the list"""
    if tree is None:
        tree = []

    try:
        terms = [_list.pop(0), _list.pop(0)]
    except IndexError:
        return tree

    frequency_total = 0
    letters = ''

    for term in terms:
        _frequency, _letters = term
        frequency_total += _frequency
        letters += _letters

    compressed = (frequency_total, ''.join(sorted(letters)))
    _list.append(compressed)
    _list = sorted(_list, key=lambda x: (x[0], x[1]))

    tree.insert(0, [compressed, terms])

    return compress_list(_list, tree)


def replace(obj, key, val):
    """Recursively replace key in dict"""
    return {k: replace(val if k == key else v, key, val) 
        for k, v in obj.items()} if isinstance(obj, dict) else obj


def get_binary_tree(tree):
    """Get the binary tree"""
    binary_dict = {}

    for branch in tree:
        key, branches = branch
        key = '{0}{1}'.format(*key)
        branch_dict = {}
        for _branch in branches:
            branch_key = '{0}{1}'.format(*_branch)
            branch_dict[branch_key] = {}
        if len(binary_dict.keys()) == 0:
            binary_dict[key] = branch_dict
        binary_dict = replace(binary_dict, key, branch_dict)

    return binary_dict


def get_binary_repr(binary_tree, letter, _repr=''):
    """Get the binary presenation of a letter"""
    if len(binary_tree.keys()) == 0:
        return _repr

    if len(binary_tree.keys()) == 1:
        return get_binary_repr(list(binary_tree.values())[0], letter, _repr)

    if len(binary_tree.keys()) == 2:
        count = 0
        for k, v in binary_tree.items():
            if letter in k:
                if count == 0:
                    _repr += '0'
                else:
                    _repr += '1'
                return get_binary_repr(v, letter, _repr)
            count += 1


def run(string, letter):
    counter = Counter(string)
    frequency_list = list(zip(counter.values(), counter.keys()))
    sorted_list = sorted(frequency_list, key=lambda x: (x[0], x[1]))
    unprocessed_tree = compress_list(sorted_list)
    binary_tree = get_binary_tree(unprocessed_tree)
    print(get_binary_repr(binary_tree, letter))


if __name__ == '__main__':
    for line in fileinput.input():
        if line.startswith('#'):
            continue
        string, letter = line.split()
        print(string, letter)
        run(string, letter)