Binary Search Tree - Python

생성일: 2021년 11월 22일 오후 6:39

바이너리 서치 트리를 파이썬으로 구현했다.

• C++과 달리 주의해야 할점
  • c++에서는 재귀를 사용할 때 노드의 주소를 레퍼런스로 주어 삭제하거나 추가할 때 자동으로 자식 노드와 부모 노드가 연결되게 할 수 있지만 파이썬에서는 자동연결이 불가능 ⇒ 재귀 함수가 return 될때 부모 노드와 연결되도록 작업해야함

TreeType.py

class TreeNode:

    def __init__(self, info):
        self.info = info
        self.left = None
        self.right = None
    
class BST():

    def __init__(self):
        self.root = None
        self.order_list = []
    
    def is_empty(self):
        return (self.root == None)
    
    def count_nodes(self, tree):
        '''[1]'''
        if (tree == None):
            return 0;
        else:
            return self.count_nodes(tree.left) + self.count_nodes(tree.right) + 1

    def length_is(self):
        return self.count_nodes(self.root)

    def insert(self, item):
        '''[2]'''
        self.root = self.insert_item(self.root, item)

    def insert_item(self, node, item):
        '''[3]'''
        if(node == None):
            node = TreeNode(item)
        elif (item < node.info):
            node.left = self.insert_item(node.left, item) #부모와 자식 노드를 이렇게 연결해야 함
        else:
            node.right = self.insert_item(node.right, item)
        return node
  
    def inorder(self, node):
        '''[4]'''
        if(node != None):
            self.inorder(node.left)
            self.order_list.append(node.info)
            self.inorder(node.right)
    
    def preorder(self, node):
        '''[5]'''
        if(node != None):
            self.order_list.append(node.info)
            self.preorder(node.left)
            self.preorder(node.right)

    def postorder(self, node):
        '''[6]'''
        if(node != None):
            self.postorder(node.left)
            self.postorder(node.right)
            self.order_list.append(node.info)

    def delete(self, item):
        '''[7]'''
        self.root, deleted = self.delete_node(self.root, item) #delete가 호출될때마다 root를 계속 갱신
        return deleted

    def delete_node(self, current, item):
        '''[8]'''
        if current is None:
            return current, False
        deleted = False
        if item == current.info: #지울 노드 찾았을 때
            deleted = True
            if current.left and current.right: #지우려는 노드 양쪽이 자식이 있을 때
                data = None
                data = BST.get_predecessor(current.left, data)#static method 사용법(class이름.함수이름)
                current.info = data
                current.left, deleted = self.delete_node(current.left, data)
            elif current.left or current.right: #자식이 한개만 있는 노드를 지울 때
                current = current.left or current.right
            else: #자식이 양쪽 다 없을 때
                current = None
        elif item < current.info:
            current.left, deleted = self.delete_node(current.left, item)
        else:
            current.right, deleted = self.delete_node(current.right, item)
        return current, deleted

    #사용해야함(static)
    def get_predecessor(tree, data):
        '''[9]'''
        while(tree.right != None):
            tree = tree.right
        data = tree.info
        return data

testTreeType.py

from TreeType import *

if __name__ == '__main__':
    tree = BST()

    tree.insert(5)
    tree.insert(2)
    tree.insert(3)
    tree.insert(1)
    tree.insert(9)
    tree.insert(8)
    tree.insert(7)

    print("Inorder Result")
    tree.inorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print("\nPreorder Result")
    tree.preorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print("\nPostorder Result")
    tree.postorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print()

    print("Delete Result")
    tree.delete(1)
    tree.delete(8)
    tree.inorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    
    print("Length is: ",tree.length_is())
    print()

ITreeType.py (BST를 재귀가 아닌 반복문을 사용하여 구현)

• 주의해야할 점
  • 부모노드와 자식 노드를 계속 변수로 트래킹하여 연결해줘야 함
  • Tree Traversal 을 반복분으로 구현하기 위해 list를 Stack처럼 사용함

from TreeType import TreeNode

class NodeType:
    def __init__(self, data):
        self.info = data
        self.left = None
        self.right = None

class IterBST():
    def __init__(self):
        self.root = None
        self.order_list = []

    def insert(self, data):
        '''[10]'''
        newNode = TreeNode(data)
        parentPtr, nodePtr, found = self.find_node(self.root, data)

        if (parentPtr is None):
            self.root = newNode
        elif (data < parentPtr.info):
            parentPtr.left = newNode #이런식으로 부모 노드와 연결
        else:
            parentPtr.right = newNode

    def find(self, key):
        '''[11]'''
        nodePtr = self.root
        found = False
        while(nodePtr is not None and not found):
            if(key == nodePtr.info):
                found = True
            elif (key < nodePtr.info):
                nodePtr = nodePtr.left
            else:
                nodePtr = nodePtr.right
        return found      

    def find_node(self, root, key):
        '''[12]'''
        nodePtr = root
        parentPtr = None
        found = False
        while(nodePtr is not None and not found):
            if (key < nodePtr.info):
                parentPtr = nodePtr
                nodePtr = nodePtr.left
            elif (key > nodePtr.info):
                parentPtr = nodePtr
                nodePtr = nodePtr.right
            else:
                found = True
        return parentPtr, nodePtr, found

    def delete(self, key):
        '''[13]'''
        self.delete_node(self.root, key)
        
    def delete_node(self, node, key):
        '''[14]'''
        parent = None
        curr = node
        while(curr is not None and curr.info != key):
            parent = curr
            if(key < curr.info):
                curr = curr.left
            else:
                curr = curr.right

        if(curr == None):
            return self.root
        
        if curr.left is None and curr.right is None: #자식이 없는 노드를 지울 때
            if(curr != self.root):
                if (parent.left == curr):
                    parent.left = None
                else:
                    parent.right = None
            else:
                self.root = None
        
        elif curr.left and curr.right: #자식이 2개인 노드를 지울 때
            data = None
            data = IterBST.get_predecessor(curr.left, data)
            curr.info = data
            self.delete_node(curr.left, data)
        
        else: #자식이 한개 있는 노드를 지울 때, 루트를 지울 때를 염두해야 함
            if curr.left:
                child = curr.left
            else:
                child = curr.right
            if (curr != self.root): #부모와 연결
                if(parent.left == curr):
                    curr = child
                    parent.left = child
                else:
                    curr = child
                    parent.right = child    
            else:
                root = child
                
        
    def inorder(self, node):
        '''[15]'''
        curr = node
        stack = []
        while True:
            if curr is not None:
                stack.append(curr)
                curr = curr.left
            
            elif(stack):
                curr = stack.pop()
                self.order_list.append(curr.info)
                curr = curr.right
            else:
                break

    def preorder(self, node):
        '''[16]'''
        if node is None:
            return        
        stack = []
        stack.append(node)
        while(len(stack) > 0):
            node = stack.pop()
            self.order_list.append(node.info)

            if node.right is not None:
                stack.append(node.right)
            if node.left is not None:
                stack.append(node.left)

    def postorder(self, node):
        '''[17]'''
        if node is None:
            return
        
        stack = []

        while(True):
            while(node):
                if node.right is not None:
                    stack.append(node.right)
                stack.append(node)
                node = node.left
            node = stack.pop()

            last = None
            if len(stack) > 0:
                last = stack[-1]

            if(node.right is not None and last == node.right):
                stack.pop()
                stack.append(node)
                node = node.right
            else:
                self.order_list.append(node.info)
                node = None
            
            if (len(stack) <= 0):
                break
                

    def get_predecessor(tree, data):
        '''[18]'''
        while(tree.right != None):
            tree = tree.right
        data = tree.info
        return data

testITreeType.py

from ITreeType import IterBST

if __name__ == '__main__':

    tree = IterBST()
    tree.insert(5)
    tree.insert(2)
    tree.insert(3)
    tree.insert(1)
    tree.insert(9)
    tree.insert(8)
    tree.insert(7)
 
    print("Inorder Result")
    tree.inorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print("\nPreorder Result")
    tree.preorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print("\nPostorder Result")
    tree.postorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()
    print()

    print("Delete Result")
    tree.delete(1)
    tree.delete(8)
    tree.inorder(tree.root)
    print(tree.order_list)
    tree.order_list.clear()