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

파이썬으로 트리의 모든 노드 합계 구하기: 재귀 함수 활용 예제

트리(Tree) 자료구조에 저장된 모든 노드의 값을 더한 합계를 구해야 하는 경우가 종종 있습니다. 이 글에서는 파이썬으로 트리 구조를 직접 만들고, 재귀 호출을 통해 모든 노드의 합을 계산하는 방법을 단계별로 살펴봅니다.

여기서는 'Tree_structure'라는 클래스를 정의하고, 루트 값 설정, 자식 노드 추가, 특정 값 검색, 전체 노드 합계 계산 등의 메서드를 구현합니다. 또한 사용자가 메뉴에서 원하는 작업을 선택하면 그에 맞는 연산이 트리에 수행되는 대화형 프로그램 형태로 작성했습니다.

예제 코드

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

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

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

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

    def summation_nodes(self):
        sum_val = self.key
        for child in self.children:
            sum_val += child.summation_nodes()
        return sum_val


tree = None

print('메뉴 (중복된 키는 허용되지 않습니다)')
print('add <데이터> at root          : 루트에 노드 추가')
print('add <데이터> below <데이터>   : 특정 노드 아래에 추가')
print('summation                     : 모든 노드의 합계 출력')
print('quit                          : 프로그램 종료')

while True:
    my_input = input('무엇을 하시겠습니까? ').split()

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        new_node = Tree_structure(data)
        sub_op = my_input[2].strip().lower()
        if sub_op == 'at':
            tree = new_node
        elif sub_op == 'below':
            key = int(my_input[3])
            ref_node = None
            if tree is not None:
                ref_node = tree.search_val(key)
            if ref_node is None:
                print('해당 키가 존재하지 않습니다.')
                continue
            ref_node.add_values(new_node)

    elif operation == 'summation':
        if tree is None:
            print('트리가 비어 있습니다.')
        else:
            summation_val = tree.summation_nodes()
            print('모든 노드의 합은 : {}'.format(summation_val))

    elif operation == 'quit':
        break

실행 결과

메뉴 (중복된 키는 허용되지 않습니다)
add <데이터> at root
add <데이터> below <데이터>
summation
quit
무엇을 하시겠습니까? add 56 at root
무엇을 하시겠습니까? add 45 below 56
무엇을 하시겠습니까? add 23 below 56
무엇을 하시겠습니까? summation
모든 노드의 합은 : 124
무엇을 하시겠습니까?

위 실행 결과에서 루트에 56을 추가한 뒤, 그 아래에 45와 23을 각각 추가했습니다. 따라서 전체 합계는 56 + 45 + 23 = 124가 됩니다.

코드 설명

  • Tree_structure 클래스: 트리의 각 노드를 표현하는 클래스입니다.

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

  • set_root: 트리의 루트(root) 값을 설정하는 메서드입니다.

  • add_values: 트리에 새로운 노드(요소)를 자식으로 추가하는 메서드입니다.

  • search_val: 트리에서 특정 키를 가진 노드를 찾는 메서드로, 자식 노드들을 순회하며 재귀적으로 탐색합니다.

  • summation_nodes: 트리의 모든 노드 값을 합산하는 메서드입니다. 자기 자신의 값을 더한 뒤, 각 자식의 합계를 재귀적으로 호출하여 누적합니다.

  • 사용자에게는 '루트에 추가(at root)', '특정 노드 아래에 추가(below)', '합계 계산(summation)', '종료(quit)' 네 가지 옵션이 제공됩니다.

  • 사용자가 선택한 옵션에 따라 해당 연산이 수행되고, 그 결과가 콘솔에 출력됩니다.

참고 사항

원본 코드의 search_val 메서드 내부에는 존재하지 않는 search 메서드를 호출하는 오타가 있어, 위 예제에서는 올바르게 동작하도록 search_val로 수정했습니다.

또한 이 방식의 합계 계산은 트리의 모든 노드를 한 번씩 방문하므로, 시간 복잡도는 노드 개수에 비례하는 O(n)입니다. 트리가 매우 깊어지면 재귀 호출 깊이 제한(Python 기본 약 1,000)에 도달할 수 있으므로, 필요하다면 sys.setrecursionlimit()으로 한도를 조정하거나 스택 기반 반복문으로 구현할 수 있습니다.