백준 15681 [트리와 쿼리] 파이썬

https://www.acmicpc.net/problem/15681

 

# 백준 15681 트리와 쿼리

import sys
sys.setrecursionlimit(10**6)

#초기화 
N,R,Q=map(int,sys.stdin.readline().split()) #입력값 받음
Prev_Tree=[[] for _ in range(N+1)] # 순서 없이 저장할 Tree
Prev_Visted=[0]*(N+1) # 순서 정하기 전 방문처리위한 List
Tree=[[] for _ in range(N+1)] # 순서 정하여(부모-자식 관계) 저장할 Tree
q=[]
ans=[0]*(N+1)

#양방향 연결
q.append(R)
Prev_Visted[R]=1
for _ in range(N-1):
    U,V=map(int,sys.stdin.readline().split())
    Prev_Tree[U].append(V)
    Prev_Tree[V].append(U)

# 단방향 연결을 위한 탐색 --> 트리 생성
while(q):
    node=q.pop(0)
    for neighborhood in Prev_Tree[node]:
        if Prev_Visted[neighborhood]==0:
            Tree[node].append(neighborhood)
            Prev_Visted[neighborhood]=1
            q.append(neighborhood)

# 트리 순회
def visit_child(node):
    cnt=1 
    for child in Tree[node]:
        cnt+=visit_child(child)
    ans[node]=cnt
    return cnt


visit_child(R)

for _ in range(Q):
    print(ans[int(sys.stdin.readline())])

문제를 부분적으로 나누어 풀었다. 

 

트리 순회 파트를 한번에 계산하지 않고, 입력이 들어왔을때마다 계산하는 방식으로 했더니 시간초과가 발생했었다.

(아래 해당코드)

for _ in range(Q):
    ans=0
    query_node=int(sys.stdin.readline())
    q.append(query_node)
    while(q):
        node=q.pop(0)
        ans+=1
        for child in Tree[node]:
            q.append(child)
    print(ans)

 

 

 

 

 

다른 분의 풀이를 첨부한다. 한번에 계산해도 되는것이었다....! (안타까운 일입니다.)

 

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])

https://velog.io/@zjvlzld/%EC%95%8C%EA%B3%A0%EB%A6%AC%EC%A6%98%EB%B0%B1%EC%A4%80-15681

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

<출처>