[백준] 2668 - 숫자고르기 [Python(파이썬)]

문제

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

2668번: 숫자고르기

 

풀이

예시의 표

윗줄에서 1, 3, 5를 뽑으면 첫째 줄과 둘째 줄에서 동일한 집합을 만들어낼 수 있다. 여기서 생각해볼 수 있는 것은 결국 다른 것들은 고려할 필요없이 첫째 줄과 둘째 줄에서 같은 열에 놓이는 대응되는 숫자들이 중요하다는 것이다! 여기서 좀 더 추상화를 해보자.

 

위의 표를 둘째 줄의 노드가 윗줄의 노드를 가리키고 있는 한 방향 그래프로 바꾸어 생각해보면, 다음과 같은 연결리스트로 표현할 수 있다.

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

이를 그림으로 표현하면 다음과 같다. 

뽑을 수 있는 1, 3, 5 모두 사이클을 형성하는 것을 확인할 수 있다. 그렇다, 이 문제는 사이클을 찾아내어, 그것을 구성하는 노드를 모두 찾아내면 되는 문제였다.

 

이제 구현으로 가보자. 중복적인 탐색을 피하기 위해 checked 리스트를 만들어 탐색 여부를 확인한다. DFS 함수 내에선 방문한 노드를 visited 집합에 담아준다. visited를 그대로 전달하면 참조가 같아지기 때문에, .copy( ) 메소드를 이용해서 새로운 visited를 만들어서 전달한다.

 

인접한 노드를 확인할 때, 이미 방문한 노드가 또 다시 등장한다는 것은 사이클이 존재한다는 것이다. 이렇게 사이클이 생기면 result 리스트에 사이클을 구성하는 노드들을 모두 추가하고, 문제에서 요구하는대로 출력하면 된다.

 

코드

import sys
from collections import defaultdict
sys.setrecursionlimit(10**6)

def dfs(u, visited):
    visited.add(u)
    checked[u] = 1
    for v in g[u]:
        if v not in visited:
            dfs(v, visited.copy())
        else: # 사이클이 생기면 뽑는다.
            result.extend(list(visited))
            return

n = int(sys.stdin.readline().strip())
g = defaultdict(list)
for i in range(1, n+1):
    v = int(sys.stdin.readline().strip())
    g[v].append(i)

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

result.sort()
print(len(result)) 
for x in result:
    print(x)