[좌충우돌, 파이썬으로 자료구조 구현하기] 이진트리 표현하기 (연결리스트)

이진트리

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

class Tree:
    def __init__(self,node=None):
        self.root = node

if __name__ == '__main__':
    tree = Tree(Node("A"))
    tree.root.left = Node("B")
    tree.root.right = Node("C")
    tree.root.left.left = Node("D")
    tree.root.left.right = Node("E")
    tree.root.right.left = Node("F")
    tree.root.left.left.left = Node('G')

위 이진트리 그림을 토대로 노드와 트리 클래스를 정의해보았다.

아래 그림은 Python Tutor를 통해 나온 결과이다.

재귀 순회 함수 만들기

배열로 표현한 이진트리를 순회할 때에는 배열의 인덱스가 트리를 벗어났는지와 원소가 None 값인지를 확인했다면,

연결리스트로 표현한 이진트리는 노드가 None인지만 검사하면 된다.

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

class Tree:
    def __init__(self,node=None):
        self.root = node

    def preOrder(self):
        def _preOrder(node):
            if node is None:
                return
            res.append(node.data)
            _preOrder(node.left)
            _preOrder(node.right)

        res = []
        _preOrder(self.root)
        return res

    def inOrder(self):
        def _inOrder(node):
            if node is None:
                return
            _inOrder(node.left)
            res.append(node.data)
            _inOrder(node.right)

        res = []
        _inOrder(self.root)
        return res

    def postOrder(self):
        def _postOrder(node):
            if node is None:
                return
            _postOrder(node.left)
            _postOrder(node.right)
            res.append(node.data)

        res = []
        _postOrder(self.root)
        return res

if __name__ == '__main__':
    tree = Tree(Node("A"))
    tree.root.left = Node("B")
    tree.root.right = Node("C")
    tree.root.left.left = Node("D")
    tree.root.left.right = Node("E")
    tree.root.right.left = Node("F")
    tree.root.left.left.left = Node('G')

    print("전위 순회 결과: ", tree.preOrder())
    print("중위 순회 결과: ", tree.inOrder())
    print("후위 순회 결과: ", tree.postOrder())

레벨 순서 순회 함수 만들기

앞서 레벨 순서 순회 함수는 queue를 활용하여 구현했었다. 참고하여 구현해보자.

def levelOrder(self):
        if not self.root:
            return []

        res,queue = [],[self.root]

        while queue:
            node = queue.pop(0)
            res.append(node.data)
            if node.left:
                queue.append(node.left)
            if node.right:
                queue.append(node.right)
        return res

배열로 이진 트리 구성하기

지금까지는 이진 트리를 구성하기 위해서 tree.root.left = Node('A') 이렇게 직접 입력해주었다면, 이제는 배열을 매개변수로 전달하고 해당 배열을 이진 트리를 반환하도록 구현해보자.

배열로 표현한 이진트리 = ["A", "B", "C", "D", "E", "F", None, "G"]

레벨 순서로 순회한 결과 = ["A", "B", "C", "D", "E", "F", "G"]

즉, 레벨 순서로 순회하면서 배열의 원소를 노드로 추가해주면 된다.

def make_tree(self,arr):
        if not arr:
            return []

        self.root = Node(arr[0])
        q = [self.root]
        index = 1 # q에 root 노드가 1개 추가된 상태이므로 인덱스는 1부터 시작

        while q and index < len(arr):
            node = q.pop(0)
            if index < len(arr) and arr[index] is not None:
                node.left = Node(arr[index])
                q.append(node.left)
            index += 1
            if index < len(arr) and arr[index] is not None:
                node.right = Node(arr[index])
                q.append(node.right)
            index += 1

• q가 비어있지 않고, index가 배열의 길이를 벗어나지 않을 때까지 반복한다.
  • q가 비어있는지 체크해주는 이유는 주어진 트리 arr가 None 요소를 많이 포함한 경우를 고려해주기 위해서이다.
• q에 있는 노드를 빼내어 해당 노드의 left, right를 설정해준다.
  • 반복문 조건에 index가 배열 길이를 벗어나지 않는지 체크하는 조건이 포함되어 있지만, left, right를 설정할 때도 해당 조건을 확인해줘야하는 이유는 반복문 내부에서 index를 추가해주고 있어 left가 추가될 때, right가 추가될 때도 확인을 해야지만 오류가 발생하지 않는다.

이진 트리 클래스 전체 코드

class Tree:
    def __init__(self,node=None):
        self.root = node

    def make_tree(self,arr):
        if not arr:
            return []

        self.root = Node(arr[0])
        q = [self.root]
        index = 1 # q에 root 노드가 1개 추가된 상태이므로 인덱스는 1부터 시작

        while q and index < len(arr):
            node = q.pop(0)
            if index < len(arr) and arr[index] is not None:
                node.left = Node(arr[index])
                q.append(node.left)
            index += 1
            if index < len(arr) and arr[index] is not None:
                node.right = Node(arr[index])
                q.append(node.right)
            index += 1

    def preOrder(self):
        def _preOrder(node):
            if node is None:
                return
            res.append(node.data)
            _preOrder(node.left)
            _preOrder(node.right)

        res = []
        _preOrder(self.root)
        return res

    def inOrder(self):
        def _inOrder(node):
            if node is None:
                return
            _inOrder(node.left)
            res.append(node.data)
            _inOrder(node.right)

        res = []
        _inOrder(self.root)
        return res

    def postOrder(self):
        def _postOrder(node):
            if node is None:
                return
            _postOrder(node.left)
            _postOrder(node.right)
            res.append(node.data)

        res = []
        _postOrder(self.root)
        return res

    def levelOrder(self):
        if not self.root:
            return []

        res,queue = [],[self.root]

        while queue:
            node = queue.pop(0)
            res.append(node.data)
            if node.left:
                queue.append(node.left)
            if node.right:
                queue.append(node.right)
        return res

    def preOrderByStack(self):
        if not self.root:
            return []

        res, stack = [], [self.root]

        while stack:
            node = stack.pop()
            res.append(node.data)
            if node.right:
                stack.append(node.right)
            if node.left:
                stack.append(node.left)
        return res

    def inOrderByStack(self):
        if not self.root:
            return []

        node = self.root
        res,stack = [], []

        while True:
            if node:
                stack.append(node)
                node = node.left
            elif stack:
                node = stack.pop()
                res.append(node.data)
                node = node.right
            else:
                break
        return res

    def postOrderByStack(self):
        if not self.root:
            return []

        res, stack = [], [self.root]
        while stack:
            node = stack.pop()
            res.append(node.data)
            if node.left:
                stack.append(node.left)
            if node.right:
                stack.append(node.right)
        return res[::-1]

• Stack으로 구현한 전위, 중위, 후외 순회 메서드로 추가로 구현해보았다.