코딩테스트 파이썬/탐색

11725 트리의 부모 찾기 파이썬

백엔드 개발자 2021. 8. 5. 12:32

문제

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다.

출력

첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

 

 

그래프문제와 같다고 생각했다. 트리는 특수한 형태의 그래프이기 때문이다.

 

입력 및 변수로

1. 노드의 개수 N

2. 그래프 배열 graph

3. 함수 -> 부모노드 찾기

 

처음에 생각한 알고리즘은 

노드마다 한번씩 함수를 돌려서

시작노드에서 pop한 값이 1개면 바로 반환.

2개면 둘중 하나만 dfs로 해서 부모노드인것을 반환하는 형식으로 구했다.

그래서 함수를 2개 썼다. 하나는 부모노드 반환, 하나는 dfs

 

결과는 시간초과였다.

 

입출력 부분 :

 

import sys
input=sys.stdin.readline
#N 입력
N=int(input())
#그래프 입력
graph=[ [] for _ in range(N+1) ]

for i in range(N-1):
    n1
,n2 =map(int,input().split())
    graph[n1].append(n2)
    graph[n2].append(n1)

 

이부분은 동일했다. graph를 빈배열로 선언해서 2중 리스트를 만들고, 

2개를 번갈아서 전부 append해주면 된다.

 

처음 풀이 

def dfs(start):
    need_visit=[start]
    visited=[]
   
while need_visit:
        node = need_visit.pop(
0)
        
if node==1:
           
return True
        if
node not in visited:
           
for i in graph[node]:
                need_visit.append(i)
            visited.append(node)
   
return False
def
FindRoot(node):
    stack=[]
   
if len(graph[node])==1:
        
return graph[node].pop()
   
else:
       
for i in graph[node]:
            stack.append(i)
       
if 1 in stack:
           
return 1
       
else:
            left=stack.pop()
            right=stack.pop()
           
if dfs(left)==True:
                
return left
           
else:
               
return right

 

dfs는 전형적인 함수구현이었고, FindRoot의 경우

graph내부원소가 1개일 때 먼저 구하고

2개일 때는 1있으면 바로 리턴

없으면 두개를 나눠서 하나만 dfs를 돌려보고

거기가 부모노드가 아니면 반대를 리턴한다.

테케는 맞게 나왔으나 시간초과가 떴다. 왜?

V+E는 최대 20만정도인데, 이걸 마지막에 N개마다 돌려서 그렇게 된것 같다.

if not in 코드가 시간복잡도가 O(N)이라고 한다..

 

 

 

참고 :https://jainn.tistory.com/26

 

이분의 풀이를 참고했다.

 

나는 각각 노드별로 생각한 반면에

한번에 dfs에서 방문여부 배열개념을 추가해서 부모노드가 아닌것들을 걸러주고,

dfs 순회를 하면서 그때그때 부모노드를 구해주는 방법을 쓰셨다.

 

 

queue=[1]
visit=[
0]*(N+1)
result={}

while queue:
    node=queue.pop()
   
for i in graph[node]:
       
if visit[i]==0:
            result[i]=node
            visit[i]=
1
           
queue.append(i)

 아주 간단하다. 단지 방문을 안한 자식노드만, 자식노드를 키로 해서 딕셔너리에 값으로 현재 추출한 부모노드를

넣는다. 왜냐하면 루트노드인 1부터 시작이기 때문에.

방문 배열로 어떻게 부모노드를 구할 수 있냐면

 

위에서 아래로 내려오기 때문에 방문한 지점은 당근 부모노드들이 된다.

그러니 자식노드에 연결된 노드들 중 방문을 이미한데는 제외가 된다.

 

여튼 결론은 위에서 아래로 내려갈 때 한번에 부모노드를 같이 구한다는 것이다.

 

 

import sys
input=sys.stdin.readline
#N 입력
N=int(input())
#그래프 입력
graph=[ [] for _ in range(N+1) ]

for i in range(N-1):
    n1
,n2 =map(int,input().split())
    graph[n1].append(n2)
    graph[n2].append(n1)

queue=[
1]
visit=[
0]*(N+1)
result={}

while queue:
    node=queue.pop()
   
for i in graph[node]:
       
if visit[i]==0:
            result[i]=node
            visit[i]=
1
           
queue.append(i)

for i in range(2, N+1):
   
print(result[i])