[Python/BOJ] 2668. 숫자고르기

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

문제

세로 두 줄, 가로로 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)을 나타내는 정수 하나가 주어진다. 그 다음 줄부터는 표의 둘째 줄에 들어가는 정수들이 순서대로 한 줄에 하나씩 입력된다.

출력

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

풀이

Code

import sys
from collections import deque
input = sys.stdin.readline

# 1. bfs 함수 정의
def bfs(node) :
    global ans

    # 1-1. 경로 변수 생성
    route = set([node])
    # 1-2. 큐 생성 후 현재 위치 삽입
    queue = deque([node])
    # 1-3.
    while queue :
        # 1-3-1. 위치 인덱스 반환
        now = queue.popleft()
        # 1-3-2. 도착지 생성
        next = array[now]
        # 1-3-3. 이미 온 적이 있는 경우
        if next in route :
            # 해당 위치가 출발지일 경우(사이클이 형성되는 경우)
            if next == node :
                # 정답 집합 업데이트
                ans = ans | route
            return
        # 1-3-4. 온 적이 없는 경우
        # 큐에 현재 위치 삽입
        queue.append(next)
        # 경로 변수 업데이트
        route.add(next)

n = int(input())
array = [0] + [int(input()) for _ in range(n)]

# 2. 정답 집합 생성
ans = set()
# 3.
for i in range(1, n+1) :
    # bfs 실행
    bfs(i)
# 4. 결과 출력
print(len(ans))
for node in sorted(ans) : print(node)