Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 트리 구성하기: 노드 삽입, 삭제, 탐색 구현 방법

파이썬에서 트리(Tree) 자료구조를 직접 구성하고, 요소를 삽입하고, 삭제하며, 트리의 모든 노드를 표시해야 하는 경우가 있습니다. 이럴 때는 필요한 기능을 메서드로 가진 클래스를 정의하고, 해당 클래스의 인스턴스를 생성하여 트리의 요소에 접근하고 각종 연산을 수행하면 됩니다.

아래에서 실제 동작 과정을 살펴보겠습니다.

예제 코드

class Tree_struct:
   def __init__(self, data=None, parent=None):
      self.key = data
      self.children = []
      self.parent = parent

   def set_root(self, data):
      self.key = data

   def add_node(self, node):
      self.children.append(node)

   def search_node(self, key):
      if self.key == key:
         return self
      for child in self.children:
         temp = child.search_node(key)
         if temp is not None:
            return temp
      return None

   def remove_node(self):
      parent = self.parent
      index = parent.children.index(self)
      parent.children.remove(self)
      for child in reversed(self.children):
         parent.children.insert(index, child)
         child.parent = parent

   def bfs(self):
      queue = [self]
      while queue != []:
         popped = queue.pop(0)
         for child in popped.children:
            queue.append(child)
         print(popped.key, end=' ')

my_instance = None

print('Menu (this assumes no duplicate keys)')
print('add <data> at root')
print('add <data> below <data>')
print('remove <data>')
print('display')
print('quit')

while True:
   do = input('What would you like to do? ').split()

   operation = do[0].strip().lower()
   if operation == 'add':
      data = int(do[1])
      new_node = Tree_struct(data)
      suboperation = do[2].strip().lower()
      if suboperation == 'at':
         my_instance = new_node
      elif suboperation == 'below':
         position = do[3].strip().lower()
         key = int(position)
         ref_node = None
         if my_instance is not None:
            ref_node = my_instance.search_node(key)
         if ref_node is None:
            print('No such key.')
            continue
         new_node.parent = ref_node
         ref_node.add_node(new_node)

   elif operation == 'remove':
      data = int(do[1])
      to_remove = my_instance.search_node(data)
      if my_instance == to_remove:
         if my_instance.children == []:
            my_instance = None
         else:
            leaf = my_instance.children[0]
            while leaf.children != []:
               leaf = leaf.children[0]
            leaf.parent.children.remove_node(leaf)
            leaf.parent = None
            leaf.children = my_instance.children
            my_instance = leaf
      else:
         to_remove.remove_node()

   elif operation == 'display':
      if my_instance is not None:
         print('Breadth First Search traversal is : ', end='')
         my_instance.bfs()
         print()
      else:
         print('The tree is empty')

   elif operation == 'quit':
      break

실행 결과

Menu (this assumes no duplicate keys)
add <data> at root
add <data> below <data>
remove <data>
display
quit
What would you like to do? add 5 at root
What would you like to do? add 6 below 5
What would you like to do? add 8 below 6
What would you like to do? remove 8
What would you like to do? display
Breadth First Search traversal is : 5 6
What would you like to do? quit

코드 설명

  • 필요한 속성을 포함한 Tree_struct 클래스가 생성됩니다.

  • __init__(생성자) 함수는 자식 노드를 담을 빈 리스트를 초기화하는 역할을 합니다.

  • set_root 메서드는 트리의 루트(root) 값을 지정할 때 사용됩니다.

  • add_node 메서드는 트리에 새로운 노드를 추가하는 데 활용됩니다.

  • search_node 메서드는 재귀 호출을 통해 특정 키(key)를 가진 노드를 검색합니다.

  • remove_node 메서드는 트리에서 노드를 삭제하며, 삭제된 노드의 자식들은 부모 노드 아래로 다시 연결됩니다.

  • bfs 메서드는 큐(queue)를 활용해 트리 전체를 너비 우선 탐색(BFS) 방식으로 순회하고 결과를 출력합니다.

  • 이후 인스턴스를 하나 생성하고 None으로 초기화합니다.

  • 사용자로부터 수행할 작업(삽입, 삭제, 표시 등)을 입력받습니다.

  • 사용자의 선택에 따라 해당 연산이 실행됩니다.

  • 연산 결과가 콘솔 화면에 출력됩니다.

이 프로그램은 메뉴 기반으로 동작하기 때문에 루트에 노드 추가, 특정 노드 아래에 노드 추가, 노드 삭제, BFS 순회 출력 등의 작업을 대화형으로 테스트해 볼 수 있다는 장점이 있습니다.