1. ADT 작성

RBTree:
    Datas:
        __root : 트리의 루트노드
        __size : 트리의 크기

    Operations:
        insert(key) : 키 값이 key인 노드를 트리에 삽입
        delete(key) : 키 값이 key인 노드를 트리에서 제거. 성공할 경우 True, 실패할 경우 False 반환
        get(key) : 키 값이 key 인 노드를 트리에서 탐색하여 해당 노드의 데이터를 반환.
        
RBTree.__Node:
    Datas:
    	key : 노드의 키 값
        color : 노드의 색상 (0: 검은색, 1: 붉은색)
        parent : 부모노드
        left : 좌측 자식노드
        right : 우측 자식노드
        IS_NIL : NIL 노드 여부
    Operations: -

 

 

2. 구현

# 레드블랙 트리 클래스
class RBTree:
    # 노드 클래스
    class __Node:
    	# 노드 생성자
        # 기본적으로 NIL 노드로 생성된다
        def __init__(self, p=None):
            # 키값은 None, 색은 0(검은색)
            self.key = None
            self.color = 0

            # 부모노드
            self.parent = p

            # 좌측 자식노드, 우측 자식노드는 None
            self.left = None
            self.right = None

            # NIL 여부 True
            self.IS_NIL = True

    # 트리 생성자
    def __init__(self):
        # 루트노드에 NIL 노드 세팅
        self.__root = self.__Node()

        # 트리의 크기 0으로 초기화
        self.__size = 0

    # len 함수의 결과를 구현
    def __len__(self):
        # 트리의 크기 반환
        return self.__size

    # 데이터 삽입 함수
    def insert(self, key):
        # 삽입할 위치 탐색
        cursor = self.__root
        while not cursor.IS_NIL:
            # 현재 노드보다 키 값이 크다면 우측 자식노드로 이동
            if cursor.key < key:
                cursor = cursor.right
            # 현재 노드보다 키 값이 작다면 좌측 자식노드로 이동
            elif cursor.key > key:
                cursor = cursor.left
            # 이미 키 값이 존재할 경우 삽입 종료(삽입안함)
            else:
                return False

        # 삽입 위치의 NIL 노드에 key, color, 좌측 자식노드, 우측 자식노드를 세팅하고 NIL 여부를 False 로 한다
        cursor.key = key
        cursor.color = 1
        cursor.left = self.__Node(cursor)
        cursor.right = self.__Node(cursor)
        cursor.IS_NIL = False

        # 트리 크기 1 증가
        self.__size += 1

        # 루트노드의 색은 검은색
        self.__root.color = 0

        # insert_fix 를 수행하여 RB 트리의 규칙을 유지
        self.__insert_fix(cursor)

    # 삽입 이후 규칙에 맞게 트리를 수정하는 함수
    def __insert_fix(self, node):
        parent = node.parent
        # 부모노드가 존재하고 붉은색인 동안 반복
        # 루트노드는 검은색이기 떄문에 부모노드가 붉은색이면 반드시 조부모 노드가 존재
        while parent and parent.color == 1:
            # 조부모노드
            grand_parent = parent.parent

            # 부모노드와 노드가 어느쪽 자식노드인지 확인
            # True: 좌측 / False: 우측
            parent_direction = (grand_parent.left == parent)
            node_direction = (parent.left == node)

            # 삼촌노드
            uncle_node = grand_parent.right if parent_direction else grand_parent.left

            # 케이스 1. Recoloring
            # 삼촌 노드가 붉은색인 경우
            if uncle_node.color == 1:
                # 삼촌노드와 부모노드를 모두 검은색으로 변경
                uncle_node.color = parent.color = 0

                # 조부모노드를 붉은색으로 변경
                grand_parent.color = 1

                # 조부모노드를 삽입노드로 하여 fix 를 다시 수행
                node = grand_parent
                parent = node.parent
                continue

            # 케이스 2. Restructuring
            # 삼촌 노드가 검은색인 경우
            else:
                # 부모노드와 노드가 서로 다른 방향의 자식노드인 경우(LR, RL 형태)
                if node_direction != parent_direction:
                    # LR 형태인 경우 LL 형태로 변형
                    if parent_direction:
                        self.__left_rotation(parent)

                    # RL 형태인 경우 RR 형태로 변형
                    else:
                        self.__right_rotation(parent)

                    # 회전에 의해 parent 와 node 가 뒤바뀜
                    node, parent = parent, node

                # LL 형태인 경우 조부모노드에 대해 right_rotation
                if parent_direction:
                    self.__right_rotation(grand_parent)
                # RR 형태인 경우 조부모노드에 대해 left_rotation
                else:
                    self.__left_rotation(grand_parent)

                # 부모노드가 된 parent 노드의 색을 검은색으로, 나머지 두 노드는 붉은색으로 한다
                parent.color = 0
                grand_parent.color = 1
                node.color = 1

                # Restructuring 은 추가작업을 필요로하지 않음
                break

        # 루트노드는 항상 검은색
        self.__root.color = 0

    # 좌측 회전 함수
    def __left_rotation(self, node):
        if node.IS_NIL or node.right.IS_NIL:
            return

        parent = node.parent
        right_tmp = node.right

        node.right = right_tmp.left
        right_tmp.left.parent = node

        right_tmp.left = node
        node.parent = right_tmp

        right_tmp.parent = parent
        if parent:
            if node == parent.left:
                parent.left = right_tmp
            else:
                parent.right = right_tmp
        else:
            self.__root = right_tmp

    # 우측 회전 함수
    def __right_rotation(self, node):
        if node.IS_NIL or node.left.IS_NIL:
            return

        parent = node.parent
        left_tmp = node.left

        node.left = left_tmp.right
        left_tmp.right.parent = node

        left_tmp.right = node
        node.parent = left_tmp

        left_tmp.parent = parent
        if parent:
            if node == parent.left:
                parent.left = left_tmp
            else:
                parent.right = left_tmp
        else:
            self.__root = left_tmp

    # 데이터 삭제 함수
    def delete(self, key):
        # 삭제할 노드 탐색
        target = self.__search(key)

        # 삭제할 노드를 찾지 못한 경우
        if not target:
            return False

        # 삭제 후 fix 를 위한 변수
        child, s_color = None, None

        # target 이 리프노드인 경우
        if target.left.IS_NIL and target.right.IS_NIL:
            child, s_color = target, target.color
            target.key = None
            target.left = target.right = None
            target.IS_NIL = True
            target.color = 0

        # 좌측 자식노드를 가진 경우
        elif not target.left.IS_NIL:
            # 좌측 서브트리중 가장 오른쪽의 노드를 구한다
            successor = target.left

            while not successor.right.IS_NIL:
                successor = successor.right

            # successor 의 좌측 자식노드와 색상 저장
            child, s_color = successor.left, successor.color

            # 삭제할 노드의 키값을 successor 의 키값으로 변경
            target.key = successor.key

            # successor 가 target 의 좌측 자식노드인 경우
            if successor == target.left:
                # successor 의 부모노드의 좌측 자식노드를 successor 의 좌측 자식노드로 설정
                successor.parent.left = successor.left

            # successor 가 target 의 좌측 자식노드가 아닌 경우
            else:
                # successor 의 부모노드의 우측 자식노드를 successor 의 좌측 자식노드로 설정
                successor.parent.right = successor.left

            # successor 의 좌측 자식노드의 부모노드를 successor 의 부모노드로 설정
            successor.left.parent = successor.parent

        # 우측 자식노드만 가진 경우
        else:
            # 우측 서브트리중 가장 왼쪽의 노드를 구한다
            successor = target.right
            while not successor.left.IS_NIL:
                successor = successor.left

            # successor 의 우측 자식노드와 색상 저장
            child, color = successor.right, successor.color

            # 삭제할 노드의 키값을 successor 의 키값으로 변경
            target.key = successor.key

            # successor 가 target 의 우측 자식노드인 경우
            if successor == target.right:
                # successor 의 부모노드의 우측 자식노드를 successor 의 우측 자식노드로 설정
                successor.parent.right = successor.right

            # successor 가 target 의 우측 자식노드가 아닌 경우
            else:
                # successor 의 부모노드의 좌측 자식노드를 successor 의 우측 자식노드로 설정
                successor.parent.left = successor.right

            # successor 의 우측 자식노드의 부모노드를 successor 의 부모노드로 설정
            successor.right.parent = successor.parent

        # 트리 크기 1 감소
        self.__size -= 1

        # 삭제이후 fix 과정
        # 트리가 비어있지 않은 경우
        if child.parent:
            # successor 가 검은색이었던 경우
            if s_color == 0:
                # child 는 붉은색인 경우
                if child.color == 1:
                    child.color = 0

                # child 도 검은색인 경우
                else:
                    self.__delete_fix(child)
        return True

    # 삭제 이후 규칙에 맞게 트리를 수정하는 함수
    def __delete_fix(self, child):
        # child 의 부모 노드
        parent = child.parent

        while parent:
            # True 일 경우 child 는 parent 의 좌측 자식노드 False 일 경우 우측 자식노드
            child_direction = (child == parent.left)

            # 형제노드
            sibling = parent.right if child_direction else parent.left

            # 조카노드
            n1, n2 = sibling.left, sibling.right

            # parent 가 붉은색인 경우
            if parent.color == 1:
                # 케이스 1. sibling, n1, n2 가 모두 검은색인 경우
                if sibling.color == n1.color == n2.color == 0:
                    # parent 와 sibling 의 색을 교환하고 종료
                    parent.color, sibling.color = sibling.color, parent.color
                    break

            # parent 가 검은색인 경우
            else:
                # 케이스 2. sibling, n1, n2도 모두 검은색인 경우
                if sibling.color == n1.color == n2.color == 0:
                    # sibling 의 색을 붉은색으로 바꾼 뒤 child 를 parent 로,
                    # parent 를 parent 의 부모노드로 하여 continue
                    sibling.color = 1
                    child, parent = parent, parent.parent
                    continue

                # 케이스 3. sibling 은 붉은색, n1, n2는 검은색인 경우
                elif sibling.color == 1 and n1.color == n2.color == 0:
                    # parent 와 sibling 의 색을 교환
                    parent.color, sibling.color = sibling.color, parent.color

                    # child 노드의 방향에 따라 회전 수행
                    if child_direction:
                        self.__left_rotation(parent)
                    else:
                        self.__right_rotation(parent)
                    continue

            # parent 의 색에 무관한 경우
            # sibling 의 색이 검은색인 경우
            if sibling.color == 0:
                # 가까운 조카노드가 n1, 먼 조카노드가 n2가 되도록 한다
                if not child_direction:
                    n1, n2 = n2, n1

                # 케이스 4. 먼 조카노드 n2가 붉은색인 경우
                if n2.color == 1:
                    # parent 와 sibling 의 색을 교환하고 n2의 색을 검은색으로 변경
                    parent.color, sibling.color = sibling.color, parent.color
                    n2.color = 0

                    # child 노드의 방향에 따라 회전 수행
                    if child_direction:
                        self.__left_rotation(parent)
                    else:
                        self.__right_rotation(parent)
                    break

                # 케이스 5. 가까운 조카노드 n1이 붉은색이고 먼 조카노드 n2가 검은색인 경우
                if n1.color == 1 and n2.color == 0:
                    # sibling 과 n1의 색을 교환
                    sibling.color, n1.color = n1.color, sibling.color

                    # child 노드의 방향에 따라 회전 수행
                    if child_direction:
                        self.__right_rotation(sibling)
                    else:
                        self.__left_rotation(sibling)
                    continue

        # 루트노드의 색은 항상 검은색
        self.__root.color = 0

    # 키 값이 key 인 노드의 데이터를 반환하는 함수
    # 여기서 노드에는 데이터 값이 따로 없기 때문에 키 값을 반환
    # 노드를 찾을 수 없다면 None 반환
    def get(self, key):
        target = self.__search(key)
        if target:
            return target.key
        else:
            return None

    # 키 값이 key인 노드를 반환하는 함수
    def __search(self, key):
        # 루트 노드부터 탐색 시작
        cursor = self.__root
        while not cursor.IS_NIL:
            # 현재 노드의 키값보다 key 가 더 크다면 우측 자식노드로 이동
            if cursor.key < key:
                cursor = cursor.right
            # 현재 노드의 키값보다 key 가 더 작다면 좌측 자식노드로 이동
            elif cursor.key > key:
                cursor = cursor.left
            # key 에 해당하는 노드를 찾았다면 노드반환
            else:
                return cursor

        # 찾지 못했을 경우 None 반환
        return None

    # 트리를 중위순회로 출력하는 함수. 
    # 붉은 노드가 연속으로 붙어있지 않은지 체크
    def in_order(self):
        self._dfs(self.__root, 0, 0)

    def __dfs(self, node, bcnt, pre):
        if node.color == 0:
            bcnt += 1
        if not node.left.IS_NIL:
            self._dfs(node.left, bcnt, node.color)
        print(node.key, ':', bcnt, ':color=', node.color, "invalid" if pre == node.color == 1 else "valid")
        if not node.right.IS_NIL:
            self._dfs(node.right, bcnt, node.color)
  • 최소한의 기능인 insert, delete, get 과 테스트용 inorder 함수만을 구현
  • 그 외의 함수나 멤버변수는 사용자측에선 볼 수 없도록 private 로 선언

'CS > 자료구조' 카테고리의 다른 글

#12 트리(Tree) 6 - 트라이(Trie)  (0) 2021.10.25
#11 트리(Tree) 5 - B-Tree  (0) 2021.10.18
#9 트리(Tree) 4 - 레드블랙 트리(Red-Black Tree)  (0) 2021.10.08
#8 트리(Tree) 3 - AVL 트리  (0) 2021.10.07
#7 트리(Tree) 2 - 힙(Heap)  (0) 2021.10.06

+ Recent posts