본문 바로가기

Python

[좌충우돌, 파이썬으로 자료구조 구현하기] 이진트리에 키 삽입하기

문제

[https://www.geeksforgeeks.org/insertion-in-a-binary-tree-in-level-order/?ref=lbp]

이진트리와 키가 주어질 때, 이진트리의 레벨 순서에서 가장 먼저 사용할 수 있는 위치에 키를 삽입하는 함수를 만들어라.

풀이

문제에서 "레벨 순서"에서 가장 먼저 사용할 수 있는 위치라고 했으니 레벨 순서로 푼다.

위 그림의 예시를 보면 11번의 오른쪽 자식이 레벨순서에서 가장 먼저 사용할 수 있는 위치이다.

def insert_key(tree,key):
    new_node = Node(key)
    if not tree.root:
        tree.root = new_node
    else:
        q = [tree.root]

        while q:
            node = q.pop(0)
            if node.left is None:
                node.left = new_node
                break
            else:
                q.append(node.left)
            if node.right is None:
                node.right = new_node
                break
            else:
                q.append(node.right)

실행 결과

tree = Tree()
tree.make_tree([10, 11, 9, 7, None, 15, 8])
print(tree.levelOrder())

>>> [10, 11, 9, 7, 15, 8]

insert_key(tree,12)
print(tree.levelOrder())

>>> [10, 11, 9, 7, 12, 15, 8]

insert_key(tree,100)
print(tree.levelOrder())

>>> [10, 11, 9, 7, 12, 15, 8, 100]