[백준] 1068 : 트리 (Python)

 

출처 : https://www.acmicpc.net/problem/1068

 

---

 

아이디어

이번 문제는 특수 그래프 중 하나인 트리라는 것을 직접적으로 알려줬으므로, 그래프 탐색 유형임을 알 수 있다. 트리 순회를 물어보는 요구사항을 바탕으로, DFS BFS 중에서 아무거나 선택해서 사용하면 된다. 필자는 DFS를 선택하였다.

그렇다면, 해당 요구사항을 바탕으로 프로그램 설계 시, 어떤 자료구조와 핵심 연산 중 하나인 삭제를 어떤 식으로 진행해야 하는지 살펴보자.

 

인접행렬 vs 인접리스트

예제 입력 1번 케이스를 바탕으로 인사이트를 파악할 수 있다.

• n : 5
• -1 0 0 1 1
• 삭제 대상 노드 번호 : 2

1번 케이스 기준, 인접행렬과 인접리스트 시각화

양방향 트리로써 간선이 2개가 존재한다. 이를 바탕으로, 인접행렬은 대칭성을 띄게 된다. 0의 개수가 많으므로, 희소 행렬 케이스를 염두해 두어야 한다. 희소 행렬의 경우, 메모리 효율이 떨어진다.

반면, 인접 리스트는 메모리 효율 면에서 동적으로 할당하기 때문에, 월등하다. 또한, 자식노드를 탐색할 시의 조회 연산에 대한 시간복잡도도 훨씬 빠름을 알 수 있다.

그러므로, 인접 리스트를 선택하였다.

 

서브트리 삭제 탐구

이번에는 예제 입력 2번 케이스를 바탕으로 설명을 진행한다.

• n : 5
• -1 0 0 1 1
• 삭제 대상 : 1

2번 케이스 기준, 리스트 삭제 시 한계 시각화

인접리스트는 head를 나타내는 배열의 인덱스, 해당 head에 속하는 자식 배열로 2차원의 행렬이 된다. 즉, 트리의 삭제 연산은 리스트 자체를 삭제하는 것과 배열의 원소를 삭제하는 것으로 나뉜다. 

즉, head 1을 삭제하고 싶으면 리스트 자체를 삭제시키면 되고, head 1 안에 있는 4를 삭제하고 싶으면, arr[1] 안의 원소인 4를 삭제해야한다.

이는 삭제 연산을 하는데 있어 시간이 추가된다는 이야기와 같다. 그러므로, 이 시간을 줄이기 위한 개선으로 '조회 시, 해당 서브트리를 탐색 대상에서 제외'시키는 방식을 생각해보았다.

DFS 진행 시, 삭제 대상 번호와 현재 탐색 노드 번호가 같으면 SKIP한다는 의미이다. 이렇게 되면, 삭제 연산에 추가적인 시간이 필요없어지게 되므로, 이 개선 방식을 선택하여 문제를 해결할 수 있다.

 

이 때, 주의점은 '루트 노드를 삭제할 때'이다. 트리의 최상단 노드를 삭제한다는 것은 트리를 없앤다는 것으로, 이 경우 0으로 정답처리를 해야한다. 이를 포함하여, 프로그램을 구현하면 된다.

 

정답 코드

# https://www.acmicpc.net/problem/1068
# 트리
# 그래프 탐색, DFS

input=__import__('sys').stdin.readline

def dfs(root):
    cnt=0 # 단말노드 개수
    child_flag=False # 자식 서브트리 존재 여부
    
    # 해당 서브트리 루트가 삭제 대상이 아닌경우
    if root!=delete_target:
        # 자식 서브트리 탐색
        for next_node in tree[root]:
            # 삭제 대상인 자식 서브트리가 아니고, 방문하지 않은 노드면
            if next_node!=delete_target and not visited[next_node]:
                visited[next_node]=1 # 방문처리
                child_flag=True # 해당 서브트리 루트의 자식 서브트리 존재 여부 갱신
                cnt+=dfs(next_node) # 자식 서브트리 방문 & 자식 서브트리의 단말노드 개수 받아오기

    # 자식 서브트리가 없으면 단말노드이므로 +1 리턴
    # 단말노드가 아니면, 부모에게 자신이 받아온 단말노드 개수 알려주기
    return +(not child_flag) or cnt
    #return 1 if not child_flag else cnt

n=int(input())
tree=[[] for _ in range(n)]
visited=[0]*(n)

root=0
for end,start in enumerate([*map(int,input().split())]):
    if start==-1:
        root=end
    else: # 양방향 트리 : 간선 2개
        tree[start].append(end)
        tree[end].append(start)

delete_target=int(input())

visited[root]=1
# 루트노드가 삭제 대상인 경우 단말노드 0
# 아니면, 단말노드 개수 찾기
print(dfs(root) if delete_target!=root else 0)