백준 2668번: 숫자 고르기(Python, DFS)

문제

세로 두 줄, 가로로 N개의 칸으로 이루어진 표가 있다. 첫째 줄의 각 칸에는 정수 1, 2, …, N이 차례대로 들어 있고 둘째 줄의 각 칸에는 1이상 N이하인 정수가 들어 있다. 첫째 줄에서 숫자를 적절히 뽑으면, 그 뽑힌 정수들이 이루는 집합과, 뽑힌 정수들의 바로 밑의 둘째 줄에 들어있는 정수들이 이루는 집합이 일치한다. 이러한 조건을 만족시키도록 정수들을 뽑되, 최대로 많이 뽑는 방법을 찾는 프로그램을 작성하시오. 예를 들어, N=7인 경우 아래와 같이 표가 주어졌다고 하자.

이 경우에는 첫째 줄에서 1, 3, 5를 뽑는 것이 답이다. 첫째 줄의 1, 3, 5밑에는 각각 3, 1, 5가 있으며 두 집합은 일치한다. 이때 집합의 크기는 3이다. 만약 첫째 줄에서 1과 3을 뽑으면, 이들 바로 밑에는 정수 3과 1이 있으므로 두 집합이 일치한다. 그러나, 이 경우에 뽑힌 정수의 개수는 최대가 아니므로 답이 될 수 없다.

입력

첫째 줄에는 N(1≤N≤100)을 나타내는 정수 하나가 주어진다. 그 다음 줄부터는 표의 둘째 줄에 들어가는 정수들이 순서대로 한 줄에 하나씩 입력된다.

출력

첫째 줄에 뽑힌 정수들의 개수를 출력하고, 그 다음 줄부터는 뽑힌 정수들을 작은 수부터 큰 수의 순서로 한 줄에 하나씩 출력한다.

1. 문제 설명: 백준 2668번 숫자고르기

• 1부터 N까지의 숫자가 있다.

• 각 숫자 i마다 하나의 숫자 a[i]가 대응된다. (1 ≤ a[i] ≤ N)

• 즉, “i → a[i]” 로 단방향 간선이 하나씩 있는 그래프라고 볼 수 있다.

• 이때, 어떤 숫자 집합 S를 골랐을 때,

  • 집합 S에 들어 있는 숫자들을 a[i]로 대응시킨 결과 집합도 다시 S가 되도록 하고 싶다.

• 조건을 만족하는 숫자들을 모두 골라,

  • 우선 개수를 출력하고

  • 그 다음 줄부터 오름차순으로 하나씩 출력하면 된다.

핵심은

“인덱스 집합”과 “대응된 값들의 집합”이 서로 같아지는 숫자들만 고르는 것
을 그래프/사이클 관점에서 푸는 것이다.

2. 아이디어: 단방향 그래프 + DFS로 사이클 찾기

입력 구조를 잘 보면, 각 i에 대해 정확히 하나의 a[i]가 주어진다.
그래프 관점에서는 이렇게 볼 수 있다.

• 노드: 1 ~ N

• 간선: i → a[i] (각 노드에서 나가는 간선이 1개)

이런 그래프는 흔히 functional graph 형태다.
이 그래프에서 “i → a[i] → a[a[i]] → …” 를 따라가다 보면 언젠가는

• 이미 방문했던 노드로 돌아오는 순간이 오고 → 사이클이 생긴다.

문제에서 원하는 집합은

“어떤 노드에서 시작해서 계속 따라갔을 때, 자기 자신을 포함하는 사이클에 속한 노드들”

이라 생각할 수 있다.

이 코드의 기본 전략

1. 그래프 구성

   • 인덱스 i에서 a[i]로 가는 간선을 갖는 리스트 graphs를 만든다.

2. 모든 시작점 i에 대해 DFS 수행

   • 매 시작 노드 i마다 visited, finish 배열을 새로 만든다.

   • DFS 도중,

     • 아직 방문하지 않은 노드는 계속 재귀 호출

     • 이미 방문했고(visited[v] == True), 아직 DFS가 끝나지 않은 노드(finish[v] == False)를 다시 만나면
       → 그 노드는 현재 DFS 경로 상에 있는 사이클의 일부다.

3. 사이클이 감지되면 결과에 추가

   • 해당 노드 v를 result에 넣고, count를 증가시키는 방식으로 개수를 관리한다.

visited는 “이 DFS에서 한 번이라도 들른 적 있냐?”
finish는 “이 노드에 대한 DFS 처리는 끝났냐?”를 의미한다.

• visited[v] == True & finish[v] == False 인 노드 v는
  → 현재 재귀 스택(경로) 위에 있는 노드이므로, 여기로 되돌아오면 사이클이다.

3. 시간 복잡도 분석

3.1. 그래프 구성

graphs = [[] for _ in range(N+1)]
for i in range(1,N+1):
    graphs[i].append(table[1][i-1])

• 노드가 N개이므로, 그래프 구성은 O(N).

3.2. DFS 전체 비용

for i in range(1,N+1):
    visited = [False for _ in range(N+1)]
    finish = [False for _ in range(N+1)]
    dfs(i)

• 각 i를 시작점으로 DFS를 한 번씩 호출하므로 겉보기엔 최대 O(N)번 DFS.

• 하지만 그래프 특성상(각 노드 당 간선 1개)
  각 DFS에서 방문할 수 있는 노드 수는 최대 O(N).

• 따라서 이 구현은 최악의 경우 O(N²) 정도로 볼 수 있다.

백준 2668의 제약 조건에서 N ≤ 100이므로,
O(N²) = 10,000 연산 정도면 충분히 여유 있다.

3.3. 공간 복잡도

• graphs: 크기 약 O(N)

• visited, finish, result: 모두 O(N)

• 따라서 전체 공간 복잡도는 O(N).

4. 코드 흐름 해설 (부분별 설명)

이제 코드 전체를 흐름대로 보면서, 각 부분이 무슨 역할을 하는지 정리해보자.

4.1. 입력 처리 및 테이블 구성

import sys

N = int(sys.stdin.readline().strip())
table = [[n for n in range(1,N+1)]]

row = []
for _ in range(N):
    row.append(int(sys.stdin.readline().strip()))
table.append(row)

• N: 숫자의 개수 (1부터 N까지)

• table[0]: [1, 2, ..., N] (인덱스 역할)

• table[1]: 각 인덱스 i에 대응되는 값 a[i]들을 저장하는 리스트

예를 들어 입력이

7
3
1
1
5
5
4
6

이라면

table[0] = [1, 2, 3, 4, 5, 6, 7]
table[1] = [3, 1, 1, 5, 5, 4, 6]

이런 식으로 “인덱스 → 값” 구조를 깔끔하게 평행 배열로 만들어 두는 역할이다.

4.2. 그래프 구성

graphs = [[] for _ in range(N+1)]
for i in range(1,N+1):
    graphs[i].append(table[1][i-1])

• graphs는 1-based 인덱스를 맞추기 위해 크기를 N+1로 잡았다.

• graphs[i]에는 “i에서 갈 수 있는 노드들”을 리스트로 저장한다.

• 여기서는 각 i마다 정확히 하나의 노드로만 갈 수 있으므로,
  graphs[i]에는 요소가 1개만 들어간다.

예시:

# i = 1 → a[1] = 3
graphs[1] = [3]

# i = 2 → a[2] = 1
graphs[2] = [1]
...

4.3. DFS 함수: 사이클 탐지

def dfs(node_number):
    global visited
    global finish
    global count

    visited[node_number] = True
    for v in graphs[node_number]:
        if not visited[v]:
            dfs(v)
        elif visited[v] and not finish[v]:
            if v not in result:
                count += 1
                result.append(v)
                return
            else:
                return
    finish[node_number] = True

이 부분이 알고리즘의 핵심이다.

① 방문 표시

visited[node_number] = True

• 현재 노드에 들어오자마자 visited를 True로 바꿔
  “이 DFS에서 한 번은 들렀다”는 걸 저장한다.

② 인접 노드 순회

for v in graphs[node_number]:

• 이 문제에서는 graphs[node_number]에 원소가 1개뿐이지만,
  일반적인 그래프 DFS 구조를 그대로 쓰고 있다.

③ 아직 방문하지 않은 경우 → 재귀 DFS

if not visited[v]:
    dfs(v)

• 처음 보는 노드라면 재귀적으로 계속 깊이 들어간다.

④ 방문했지만 아직 종료되지 않은 노드 → 사이클 감지

elif visited[v] and not finish[v]:
    if v not in result:
        count += 1
        result.append(v)
        return
    else:
        return

• visited[v] == True 이고 finish[v] == False 라는 것은,

  • 현재 DFS 경로(재귀 호출 스택) 위에 있는 노드라는 뜻이다.

• 이런 노드를 다시 만났다는 것은 사이클(순환)을 찾았다는 의미다.

여기서는 사이클을 발견했을 때

• result 리스트에 그 노드 v를 추가

• 전체 카운트 count를 1 증가

시키고 있다.
(DFS를 종료하기 위해 return을 사용)

⑤ DFS 끝마친 노드 표시

finish[node_number] = True

• node_number에서 출발하는 모든 탐색이 끝났다는 표시.

• 이후 다른 경로에서 이 노드를 다시 만났을 때,
  “이 노드에서 더 깊이 들어가봐야 새로운 사이클은 없다”는 정보로 사용된다.

4.4. 전체 DFS 반복 및 결과 출력

count = 0 
result = []
for i in range(1,N+1):
    visited = [False for _ in range(N+1)]
    finish = [False for _ in range(N+1)]
    dfs(i)

print(count)
result.sort()
for r in result:
    print(r)

1. count, result 초기화

   • count: 사이클에 해당하는 숫자들의 개수

   • result: 사이클에 속한다고 판단된 숫자들의 리스트

2. 모든 i를 시작점으로 DFS 실행

   for i in range(1,N+1):
       visited = [False for _ in range(N+1)]
       finish = [False for _ in range(N+1)]
       dfs(i)
   

   • 각 시작점마다 visited, finish를 새로 초기화한다.

   • 이렇게 하면, i마다 “이 i를 출발점으로 생각했을 때의 방문 상태”를 새로 세팅하는 셈이다.

3. 결과 출력

   print(count)
   result.sort()
   for r in result:
       print(r)
   

   • count를 먼저 출력하고

   • result를 오름차순 정렬해 하나씩 출력한다.

5. 전체 코드

마지막으로, 앞에서 설명한 내용을 모두 포함한 전체 코드를 정리하면 다음과 같다.

import sys

N = int(sys.stdin.readline().strip())
table = [[n for n in range(1, N+1)]]

row = []
for _ in range(N):
    row.append(int(sys.stdin.readline().strip()))
table.append(row)

graphs = [[] for _ in range(N+1)]
for i in range(1, N+1):
    graphs[i].append(table[1][i-1])

def dfs(node_number):
    global visited
    global finish
    global count

    visited[node_number] = True
    for v in graphs[node_number]:
        if not visited[v]:
            dfs(v)
        elif visited[v] and not finish[v]:
            if v not in result:
                count += 1
                result.append(v)
                return
            else:
                return
    finish[node_number] = True

count = 0
result = []
for i in range(1, N+1):
    visited = [False for _ in range(N+1)]
    finish = [False for _ in range(N+1)]
    dfs(i)

print(count)
result.sort()
for r in result:
    print(r)