[BOJ] 2668 숫자고르기 (Python3)

문제

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

1	2	3	4	5	6	7
3	1	1	5	5	4	6

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

입력

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

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

정답코드

n = int(input())
dic = {}


for i in range(1, n + 1):
    dic[i] = int(input())
n = 7
dic = {1: 3, 2: 1, 3: 1, 4: 5, 5: 5, 6: 4, 7: 6}
dic = {1 :2, 2:3, 3:4, 4:1, 5:5}


answer = []
visited = [False for _ in range(n + 1)]
        
for i in dic:
    if dic[i] == i:
        visited[i] = True
        answer.append(i)
        
def loop(startKey):
    key = startKey
    log = []
    
    while True:
        value = dic[key]
        
        if value == startKey:
            log.append(key)
            for l in log:
                visited[l] = True
            return log
        
        if key in log:
            return []
        
        log.append(key)
        key = value
        
for i in dic:
    if not visited[i]:
        answer += loop(i)
        
answer.sort()

print(len(answer))

for a in answer:
    print(a)

5
1
2
3
4
5

해설

원래는 단순하게 딕셔너리 형태로 받은다음 for문을 돌며 서로가 서로를 갖고있는 케이스만을 정답으로 하게 구현했는데, 역시 골드 5의 풀이가 이렇게 단순할리가 없었다.

그래서 바로 반례를 생각해봤는데, 다음과 같은 경우가 있을 수 있었다.

1	2	3	4
2	3	4	1

이런 경우에는 순환적으로 값을 갖고있기 때문에, 답은 [1, 2, 3, 4]가 된다. 따라서 서로가 서로를 갖지 않아도 이런식으로 동일한 집합이 나올 수 있었다.

따라서 여러케이스를 두고 생각해보니, 결국 순환된 고리가 존재하면 해당 순환이 우리가 원하는 집합의 부분집합이라는 것을 깨달았다.

예를 들어, 위의 표에서 위를 key 아래를 value라고 하자.

1
1

이 경우는 key가 곧 value이며 순환이 0인 부분집합이다.

1	2
2	1

이 경우는 1이 2의 value를 갖고, 2는 1을 value로 가지므로 순환이 1인 부분집합이다.

1	2	3
2	3	1

이 경우는 1이 2의 value를 갖고, 2는 3을 value로 갖고, 3은 1을 value로 가지므로, 순환이 2인 부분집합이다.

이렇게 주어진 표를 하나씩 방문해가며 순환고리가 존재하는 부분집합은 방문처리를 해주며 결국 순환되는 숫자들을 찾아주면 된다.