백준 2668) 숫자고르기 (python)

문제

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

풀이

아래줄의 숫자가 윗줄과 같을 때, 그래프에 넣어준다고 하면 예시 문제는 다음과 같이 표현할 수 있다.

1: [2, 3]
2: []
3: [1]
4: [6]
5: [4, 5]
6: [7]
7: []

 

이 그래프를 그림으로 그린다면,아래와 같이 표시할 수 있게 된다

.

 

정답인 1,3,5가  사이클로 돌아가는 것을 알 수 있다. 

이 문제는 그래프를 dfs탐색하여, 사이클이 생길때마다 것을 정답 리스트에 보관하여 해결할 수 있는 것이다.

풀이 코드

from collections import defaultdict

n=int(input())

dict = defaultdict(list)

def dfs(x, visited):
    visited.add(x)
    checked[x] = 1
    for v in dict[x]:
        if v not in visited:
            dfs(v, visited.copy())
        else: 
            ans.extend(list(visited))
            return
        
for i in range(1,n+1):
    v = int(input())
    dict[v].append(i)

checked = [0 for _ in range(n+1)]
ans = []

for i in range(1, n+1):
    if not checked[i]:
        dfs(i, set([]))

ans.sort()
print(len(ans)) 
for i in ans:
    print(i)