클래스를 사용하지 않고 트리 구조 나타내기

백준 11725번 트리의 부모 찾기, 1167 트리의 지름 문제를 풀었다.

정점과 간선 정보가 주어지고 이를 트리 구조로 나타내고 싶었다. 트리를 어떻게 만들지 고민했다.

 

트리 구조는 클래스 선언해서 필드에 연결된 자식을 저장하는 형식으로 종종 구현했다.

이번에는 인접리스트 형식으로 구현했다. 왜냐하면 각 노드에 인덱스로 바로 접근할 수 있는 구조를 사용하고 싶었기 때문이다.

 

N개의 노드를 가진 트리의 노드 번호는 1부터 N까지 있다고 가정했을 때, 노드 번호를 통해 자식이 누구인지 확인할 수 있게 구현했다. 노드와 간선 정보는 변수 connection에 인접리스트로 주어진다.

def make_tree(connection):
    children_of = [[] for _ in range(len(connection))]
    parent_of = [-1 for _ in range(len(connection))]
    parent = 1
    parent_of[parent] = parent
    stack = [parent]
    while stack:
        parent = stack.pop()
        for child in connection[parent]:
            if parent_of[child] == -1:
                parent_of[child] = parent
                children_of[parent].append(child)
                stack.append(child)
    return children_of


connection = [0, [3], [4], [1, 4], [2, 3, 5], [4]] # connection[i]는 i번 노드와 직접 연결된 모든 노드 번호 묶음
tree = make_tree(connection)
print(tree)