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

문제

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

첫 풀이 코드

import sys
from itertools import combinations
input = sys.stdin.readline

n = int(input())
array = [int(input()) for _ in range(n)]
numbers = [i for i in range(1,n+1)]

def find_match_set(arr):
    new_arr = []
    for i in arr:
        new_arr.append(array[i-1])
    
    arr.sort()
    new_arr.sort()

    if arr == new_arr:
        return True
    else:
        return False
    

def solution():
    answer = []
    for i in range(n+1):
        combs = combinations(numbers[0:], i)
        for comb in combs:
            if find_match_set(list(comb)) and len(answer) < len(list(comb)):              
                answer = list(comb)
    
    answer.sort()
    print(len(answer))
    for i in answer:
        print(i)

    return answer

solution()

 

n의 범위가 100이하여서 combinations을 사용해서 풀 수 있을 것 같았다.

결과는 시간초과.. 

따라서 dfs를 사용해서 다시 문제를 풀어보았다.

 

최종 코드

import sys
input = sys.stdin.readline

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

def dfs(prev,num):
    visited[prev] = True

    for i in array[prev]:
        if not visited[i]:
            dfs(i,num)
        elif visited[i] and i == num: # 사이클 생성
            answer.append(i)
    
for i in range(1, n+1):
    array[int(input())].append(i)

for i in range(1, n+1):
    visited = [False] * (n+1)
    dfs(i,i)

print(len(answer))
for i in answer:
    print(i)

 

dfs의 prev와 num의 역할은 다음과 같다.

prev = dfs의 시작 부분, dfs를 타며 계속 변하는 값

num = 한 번의 dfs에서 계속해서 비교할 number (사이클이 생성되었는지 비교)

 

이미 방문한 i값인데 사이클이 생성되었다 = 표의 위 아래 값의 집합 값이 같다이므로 answer에 append 해준다.