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

파이썬으로 트리(Tree)의 리프 노드 개수 구하기

트리 자료구조에서 리프 노드(leaf node), 즉 자식 노드가 없는 말단 노드의 개수를 세어야 하는 경우가 있습니다. 이를 위해 Tree_structure 클래스를 정의하고, 루트 노드 설정, 자식 노드 추가, 특정 값 검색 등의 메서드를 함께 구현할 수 있습니다.

사용자에게는 여러 가지 메뉴 옵션이 제공되며, 사용자가 선택한 명령에 따라 트리 요소에 대한 해당 연산이 수행됩니다.

예제 코드

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

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

    def add_vals(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(key)
            if temp is not None:
                return temp
        return None

    def count_leaf_node(self):
        leaf_nodes = []
        self.count_leaf_node_helper_fun(leaf_nodes)
        return len(leaf_nodes)

    def count_leaf_node_helper_fun(self, leaf_nodes):
        if self.children == []:
            leaf_nodes.append(self)
        else:
            for child in self.children:
                child.count_leaf_node_helper_fun(leaf_nodes)

tree = None

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

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

    operation = my_input[0].strip().lower()
    if operation == 'add':
        data = int(my_input[1])
        newNode = Tree_structure(data)
        sub_op = my_input[2].strip().lower()
        if sub_op == 'at':
            tree = newNode
        elif sub_op == 'below':
            my_pos = my_input[3].strip().lower()
            key = int(my_pos)
            ref_node = None
            if tree is not None:
                ref_node = tree.search_val(key)
            if ref_node is None:
                print('No such key.')
                continue
            ref_node.add_vals(newNode)

    elif operation == 'count':
        if tree is None:
            print('The tree is empty')
        else:
            count = tree.count_leaf_node()
            print('The number of leaf nodes are : {}'.format(count))

    elif operation == 'quit':
        break

실행 결과

Menu (this assumes no duplicate keys)
add <data> at root
add <data> below <data>
count
quit
What operation would you like to perform ? add 78 at root
What operation would you like to perform ? add 90 below 78
What operation would you like to perform ? add 8 below 78
What operation would you like to perform ? count
The number of leaf nodes are : 2
What operation would you like to perform ? quit

코드 설명

  • 먼저 Tree_structure 클래스를 생성합니다.

  • 생성자(__init__)에서는 노드의 값을 저장할 key와 자식 노드들을 담을 빈 리스트 children을 초기화합니다.

  • set_root_node 메서드는 트리의 루트(root) 값을 설정하는 역할을 합니다.

  • add_vals 메서드는 트리에 새로운 노드(요소)를 추가할 때 사용됩니다.

  • search_val 메서드는 트리 내에서 특정 키 값을 가진 노드를 재귀적으로 탐색하여 찾아줍니다.

  • count_leaf_node 메서드는 트리의 리프 노드 개수를 반환하며, 실제 탐색 작업은 헬퍼 함수에 위임합니다.

  • count_leaf_node_helper_fun은 재귀적으로 호출되는 함수로, 자식이 없는 노드(리프 노드)를 리스트에 추가합니다.

  • 메뉴에는 '루트에 추가(add at root)', '특정 노드 아래에 추가(add below)', '리프 노드 개수 세기(count)', '종료(quit)'의 네 가지 옵션이 제공됩니다.

  • 사용자가 입력한 옵션에 따라 각각의 연산이 수행되며, 그 결과가 콘솔에 출력됩니다.

위 예제에서는 루트에 78을 추가한 뒤, 그 아래에 90과 8이라는 두 개의 자식 노드를 추가했습니다. 이 시점에서 리프 노드는 90과 8 두 개이므로, count 명령 실행 시 결과로 2가 출력되는 것을 확인할 수 있습니다.