(알고리즘)백준 15681 python 파이썬

문제

먼저 문제를 살펴보자
https://www.acmicpc.net/problem/15681

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)
이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V) 이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.
이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)
입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

예제입력

9 5 3
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8
5
4
8

예제정답

9
4
1

풀이방법

이 문제는 DFS로 접근해서 풀이를 하게 되었다.
양방향 트리로 입력값에 대한 트리를 만든 후 루트노드를 받아 해당 노드부터 DFS탐색을 실행한다
해당 노드에서 방문 가능한 노드의 서브트리의 개수를 리턴받아 내 노드의 서브트리 노드에 더하는 방식으로 구현했으며 마지막 리프노드의 경우 더 이상 접근 가능한 노드가 없어 1을 리턴하게 된다.
입력 예제를 그림으로 구현하면 다음과 같다.

해당 그림에서 자신 1과 자식의 서브트리 개수를 더한 값이 내 서브트리의 값이다

코드

import sys
sys.setrecursionlimit(1000000000) #반복제한 늘리기
N,R,Q=map(int,sys.stdin.readline().split(' '))


m=[[]for _ in range(N+1)]
visit=[-1 for _ in range(N+1)]

for _ in range(N-1):
    a,b=map(int,sys.stdin.readline().split(' '))
    m[a].append(b)
    m[b].append(a)

def dfs(now):
    global visit
    visit[now]=1#나 자신을 추가해준다
    for i in m[now]:
        if visit[i]==-1: #방문하지 않은 방문 가능 노드가 있다면
            visit[now]+=dfs(i) #방문하며 그 노드의 서브트리 개수를 더해준다
    return visit[now] #내 서브트리 개수를 리턴한다
dfs(R)
for _ in range(Q):
    get=int(sys.stdin.readline())
    print(visit[get])

후기

해당 문제를 처음 제출했을때는 시간 초과가 났었다. 보다보니 입력값이 상당했다.
input()을 sys.stdin.readline()으로 바꿔주니 바로 해결이 되었다.
이런 문제들 좀 불편...
알고리즘 시험 보면 가끔 sys를 못쓰게 하는 경우가 있는데 이럴 경우 어떻게 대처해야 할지 생각해야 할 것 같다.