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

파이썬(Python)으로 트리의 모든 노드 합계 구하기 – 완전 예제 코드

트리(tree) 자료구조에서 모든 노드 값의 합을 구해야 하는 경우가 종종 있습니다. 이럴 때는 클래스를 하나 정의하고, 그 안에 루트 노드를 설정하는 메서드, 트리에 노드를 추가하는 메서드, 특정 값을 검색하는 메서드, 그리고 모든 노드를 순회하며 합계를 계산하는 메서드 등을 구현하면 됩니다. 이후 이 클래스의 인스턴스를 생성하여 각 메서드를 자유롭게 호출하고 활용할 수 있습니다.

아래는 실제 동작 과정을 보여주는 예제입니다.

예제 코드

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

    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 sum_node(self):
        my_summation = self.key
        for child in self.children:
            my_summation = my_summation + child.sum_node()
        return my_summation

my_instance = None

print('Menu (assume no duplicate keys)')
print('add <data> at root')
print('add <data> below <data>')
print('sum')
print('quit')

while True:
    my_input = input('What operation would you do ? ').split()

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        new_node = Tree_struct(data)
        suboperation = my_input[2].strip().lower()
        if suboperation == 'at':
            my_instance = new_node
        elif suboperation == 'below':
            position = my_input[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
            ref_node.add_node(new_node)

    elif operation == 'sum':
        if my_instance is None:
            print('The tree is empty')
        else:
            my_summation = my_instance.sum_node()
            print('Sum of all nodes is: {}'.format(my_summation))

    elif operation == 'quit':
        break

실행 결과

Menu (assume no duplicate keys)
add <data> at root
add <data> below <data>
sum
quit
What operation would you do ? add 5 at root
What operation would you do ? add 7 below 5
What operation would you do ? add 0 below 7
What operation would you do ? sum
Sum of all nodes is: 12
What operation would you do ? quit

코드 설명

  • 필요한 속성들을 가진 'Tree_struct' 클래스를 정의합니다.

  • '__init__' 생성자는 노드의 데이터(key)를 저장하고, 자식 노드들을 담기 위한 빈 리스트(children)를 초기화합니다.

  • 'set_root' 메서드는 트리의 루트 값을 설정하는 역할을 합니다.

  • 'add_node' 메서드는 새로운 노드를 자식 목록에 추가하여 트리를 확장합니다.

  • 'search_node' 메서드는 재귀 호출을 통해 특정 키 값을 가진 노드를 찾아 반환합니다. 찾지 못하면 None을 반환합니다.

  • 'sum_node' 메서드는 현재 노드의 값에 모든 자식 노드들의 합을 재귀적으로 더해 전체 합계를 계산합니다.

  • 클래스 외부에서 인스턴스 변수를 생성하고 None으로 초기화합니다. 아직 트리가 만들어지지 않은 상태를 의미합니다.

  • 무한 반복문 안에서 사용자로부터 수행할 작업(노드 추가, 합계 계산, 종료)을 입력받습니다.

  • 사용자의 선택에 따라 조건문이 분기되어 해당 연산이 수행되며, 존재하지 않는 키에 노드를 추가하려 하면 'No such key'라는 안내 메시지가 출력됩니다.

  • 'sum' 명령 입력 시 빈 트리 여부를 먼저 확인한 뒤, 노드가 있다면 전체 합계를 콘솔에 출력합니다.

핵심 포인트 정리

  • 합계 계산은 깊이 우선 탐색(DFS) 방식의 재귀 호출로 이루어집니다. 즉, 루트부터 시작해 각 서브트리를 순차적으로 내려가며 값을 누적합니다.

  • 노드 개수를 n이라고 할 때, 모든 노드를 한 번씩만 방문하므로 시간 복잡도는 O(n)입니다.

  • 이 구조는 자식 수에 제한이 없는 일반 트리(n-ary tree)에도 그대로 적용할 수 있어 확장성이 좋습니다.