[BOJ] 백준 2668번 - 숫자고르기 (Python)

문제

문제 바로가기

해결

두 행에서 중점적으로 봐야할 것은 idx에 해당하는 리스트의 value값을 또다시 idx로 넣어 재귀적으로 value를 가져왔을 때 결국은 초기 idx와 동일해진다는 것이다. 그 외의 숫자에 대해서는 신경을 쓰지 않아도 된다.

이를 그래프로 나타내면, 타고타고 결국 자기 자신으로 되돌아 오거나, 자기 자신을 가리키는 사이클구조가 될 것이다.

결국 이 문제에서 중점적으로 구현해야 할 요소는 그래프를 만든 뒤, 사이클을 만드는 노드들을 파악해 개수와 정보를 반환하는 것이다.

우선 문제의 그래프는 유향 그래프이기 때문에 역방향 간선의 존재 여부를 파악한다. 역방향 간선이란 DFS 스패닝 트리에서 자식 -> 부모로 거슬러 올라가는 간선이다. 해당 내용은 아래의 포스팅에서 참고하자.

Inspiration - [알고리즘] 방향 그래프에서 사이클 찾기

역방향 간선을 찾아내는 것은 visited와 finished로 구현할 수 있는데 그림으로 간단히 수행방식을 설명하면 다음과 같다.

dfs(1)을 수행하는 도중 1을 방문처리한 뒤 dfs(2)를 실행한다. 그리고 마찬가지로 dfs(2)를 수행하는 도중 2를 방문처리한 뒤 dfs(3)을 실행한다. 이때 dfs(3)은 자기 자신을 방문처리한 뒤 dfs(1)을 실행하게 되는데 이때 dfs(1)은 방문처리되었지만 함수가 끝나지는 않은 상태이므로 역방향 간선을 발견했다고 할 수 있다. 이 과정에서 방문한 모든 노드를 리스트에 추가하면 되는 것이다.

이 과정을 모든 노드에 대해 반복하는 것을 코드로 구현해보자.

# 입력
n = int(input())
graph = [[] for i in range(n + 1)]

for i in range(1, n + 1):
  graph[i].append(int(input()))


def dfs(s):
  global result, cnt

  visited[s] = True # 방문 처리
  for now in graph[s]:
    # 방문 안했을 경우 방문, 재귀
    if visited[now] == False:
      dfs(now)
    # 방문했는데 안끝났을 경우, 역방향 간선 확인
    elif visited[now] == True and finished[now] == False:
      if now not in result: # result에 중복 삽입 방지
        cnt += 1
        result.append(now)
        return
    else:
      return
  finished[s] = True # 함수 종료


result = []
cnt = 0

# 각각의 노드마다 루트 노드로서 반복해주어야 한다.
for i in range(1, n + 1):
  visited = [False] * (n + 1)
  finished = [False] * (n + 1)
  dfs(i)

result.sort()
print(cnt)
for i in result:
  print(i)