huffman-py/huffman_py/Tree.py

#
#     Sam Hadow - Huffman-py
#     Copyright (C) 2023
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

# -*- encoding: utf-8 -*-
from huffman_py.Node import *

class Tree(object):
        def __init__(self, root = None):
            # tree root (Node object)
            self.root = root

            # tree nodes (list of Node object)
            self.nodes = []
            self.gen_tree()

            self.set_id()


        # getters and setters (private attributes)

        def get_root(self):
            return self.__root
        def set_root(self, valeur):
            self.__root = valeur
        root = property(fget=get_root, fset=set_root, doc="tree root")

        def get_nodes(self):
            return self.__nodes
        def set_nodes(self, valeur):
            self.__nodes = valeur
        nodes = property(fget=get_nodes, fset=set_nodes, doc="nodes list")

        def gen_tree(self):
            # generate nodes list from a tree root
            L = list()
            L.append(self.root)
            while len(L) > 0:
                treatment = L.pop(0)
                self.nodes.append(treatment)
                if treatment.left:
                    L.append(treatment.left)
                if treatment.right:
                    L.append(treatment.right)

        def set_id(self):
            # Give unique identifiers to every node in a tree
            for i,elem in enumerate(self.nodes):
                elem.identifier = i

        def __isub__(self, node):
            # delete node (overload -=)
            if self.seek(node) != None:
                self.nodes.remove(node)
            # recursively delete children
            if node.right != None:
                self.__isub__(node.right)
            if node.left != None:
                self.__isub__(node.left)
            return self

        def __iadd__(self, tree2):
            # Tree merge
            # new root
            occurrence = self.root.occurrence
            old_root =  self.root
            if occurrence > tree2.root.occurrence:
                self.root = Node(occurrence, '', left = tree2.root, right = old_root)
            else:
                self.root = Node(occurrence, '', left = old_root, right = tree2.root)

            # add new root to nodes list
            self.nodes.append(self.root)

            # add tree2 nodes
            self.nodes += tree2.nodes

            # regenerate unique identifiers
            self.set_id()

            return self

        def seek(self, node_to_seek):
            for node in self.nodes:
                if node == node_to_seek:
                    return node
            return None