[BOJ/백준] 20924 트리의 기둥과 가지

import sys
input = sys.stdin.readline
sys.setrecursionlimit(500000)
def root_dfs(node):
    if visited[node]:
        return
    visited[node] = True
    child_list = []
    for next_node in tree[node]:
        if not visited[next_node]:
            child_list.append(next_node)

    if len(child_list)==2:
        return [node,0]
    elif len(child_list) == 1:
        return_value = root_dfs(child_list[0])
        return_value[1] += tree[node][child_list[0]]
        return return_value
    else:
        return [node,0]

def leef_dfs(node):
    if visited[node]:
        return
    visited[node] = True
    child_list = []
    for next_node in tree[node]:
        if not visited[next_node]:
            child_list.append(next_node)
    
    if len(child_list)==0:
        return 0
    else:
        result = 0
        for child_node in child_list:
            result = max(result,leef_dfs(child_node)+tree[node][child_node])
        return result
N,R = map(int,input().split())
tree = [{} for _ in range(N+1)]
for _ in range(N-1):
    x,y,b = map(int,input().split())
    tree[x][y] = b
    tree[y][x] = b


visited = [False for _ in range(N+1)]

root_node_info = root_dfs(R)
visited[root_node_info[0]] = False
leef_dis = leef_dfs(root_node_info[0])
print(root_node_info[1],leef_dis)

2개의 DFS로 나눠서 최초엔 풀었다.

 

첫번째 풀이는 두개로 나뉘어서 DFS를 했다.

 

첫 DFS는 기가노드를 찾는것이다. 기가노드를 찾으면 그때의 길이를 저장해놓는다.

 

그리고 기가노드에서 부터 리프노드까지 최대 길이를 찾아내는 방식으로 해서 구했다.

 

import sys
input = sys.stdin.readline
sys.setrecursionlimit(300000)
def root_dfs(node,dis):
    if visited[node]:
        return
    visited[node] = True
    count = 0
    nexts = -1
    distance = 0
    for next_node in tree[node]:
        if not visited[next_node]:
            count += 1
            nexts = next_node
            distance += tree[node][next_node]
    if count == 1:
        return root_dfs(nexts,dis+distance)
    else:
        return [node,dis]

def leaf_dfs(node,dis):
    global leef_dis
    visited[node] = True
    count = 0
    for next_node in tree[node]:
        if not visited[next_node]:
            leaf_dfs(next_node,dis+tree[node][next_node])
            count += 1
    if not count:
        leef_dis = max(leef_dis,dis)


N,R = map(int,input().split())
tree = [{} for _ in range(N+1)]
for _ in range(N-1):
    x,y,b = map(int,input().split())
    tree[x][y] = b
    tree[y][x] = b


visited = [False for _ in range(N+1)]
root_dis = root_dfs(R,0)
leef_dis = 0
leaf_dfs(root_dis[0],0)
print(root_dis[1],leef_dis)

두번째 풀이는 좀더 깔끔한 풀이이다.

 

count를 별도로 세주어서 기가노드인지 리프노드인지 구분을 해주었다.

 

기가노드가 아닐때에는 재귀를 해주면서 distance를 더해주고, 만약 기가노드일때는 재귀를 머추고, 기가노드의 위치와

 

지금까지 누적된 거리를 반환해준다.

 

leaf_dfs도 비슷한 과정을 겪는다.

 

leaf_dfs를 계속 재귀를 통해 가면서 leaf_node에 도착하면 현재까지 누적된 값과 최대값과 비교를 해서 더 큰 값을 넣어주면 된다.