[BOJ] 15681 - 트리와 쿼리 (Python)

 

 

아이디어

 

루트 노드에서 시작한 DFS로 리프 노드에서부터 서브 트리에 속한 노드 개수를 구해 나간다.

 

 

풀이 #1

 

간선 정보를 2차원 리스트 graph에 저장하고, 서브 트리의 노드 개수를 저장하는 N+1 길이의 리스트 dp를 만든다.

 

함수 dfs에서 현재 노드의 서브 트리에 속한 노드 개수를 저장하는 cnt를 현재 노드를 고려하여 1로 초기화한다.

for문에서 자식 노드들의 간선 정보에 포함된 현재 노드를 지우고, 자식 노드들의 반환값을 cnt에 더한다.

자식 노드들의 반환값을 모두 더한 cnt를 현재 노드 번호를 인덱스로 dp에 저장한 후 반환한다.

리프 노드에서는 for문이 실행되지 않고 처음 저장된 1이 그대로 반환된다.

 

루트 노드을 인자로 함수 dfs를 호출한 후, 쿼리에 따라 서브 트리의 노드 개수를 출력한다.

 

import sys

sys.setrecursionlimit(100002)
input = sys.stdin.readline

def dfs(idx):
    cnt = 1
    for i in graph[idx]:
        graph[i].remove(idx)
        cnt += dfs(i)
    dp[idx] = cnt
    return dp[idx]

N, R, Q = map(int, input().split())
graph = [[] for _ in range(N+1)]
dp = [0]*(N+1)
for _ in range(N-1):
    u, v = map(int, input().split())
    graph[u].append(v)
    graph[v].append(u)
dfs(R)
for _ in range(Q):
    u = int(input())
    print(dp[u])

 

 

리스트 visited를 이용하여 방문 여부를 확인하는 코드보다 remove 메서드로 연결 정보를 지우는 코드가 더 효율적이었다.

# 1.
for i in graph[idx]:
    if visited[i] == 0:
        visited[i] = 1
        cnt += dfs(i)

# 2. 메모리를 더 적게 사용하고 더 빠른 코드
for i in graph[idx]:
    graph[i].remove(idx)
    cnt += dfs(i)

 

 

 

풀이 #2

 

다른 분의 코드를 참고한 결과 재귀 호출이 아닌 while문으로 DFS를 구현하는 코드가 더 효율적이었다.

 

while문에서 DFS하면서 부모 노드와 자식 노드 번호를 튜플로 리스트 path에 저장한 다음,

while문 종료 후에 path를 역순으로 읽으며 dp에 서브 트리의 노드 개수를 저장해 나간다.

 

import sys

input = sys.stdin.readline

def solution():
    N, R, Q = map(int, input().split())
    graph = [[] for _ in range(N+1)]
    dp = [1]*(N+1)
    for _ in range(N-1):
        u, v = map(int, input().split())
        graph[u].append(v)
        graph[v].append(u)
    stack = [R]
    path = []
    while stack:
        idx = stack.pop()
        for i in graph[idx]:
            graph[i].remove(idx)
            stack.append(i)
            path.append((idx, i))
    for u, v in path[::-1]:
        dp[u] += dp[v]
    for _ in range(Q):
        u = int(input())
        print(dp[u])

solution()