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

문제로 가기!

처음에는 윗칸의 수를 고르고 그 아래 칸의 수를 보고, 각자 set에 넣어서 두 set이 같은지 다른지 판단해주고.. 같으면 또 다른 (윗칸) 수를 골라보고... 다르면 그 아래칸의 수를 윗칸 수로 하는 수를 찾고... 그런 식으로 생각을 해서 코드를 짜긴 했는데, 틀렸다.

질문 검색에서 다른 분들이 주신 반례를 보고, 그림을 그려보니까 문제가 좀 더 잘 이해되었다.

방향이 있는 그래프(윗 칸 -> 아랫 칸)로 보고, 인접 리스트로 입력을 받았다.
DFS로 탐색을 하면서 사이클이 발생한 것을 다 더하면 되는 문제였다.

나는 set을 활용해서 윗 칸의 수들의 집합, 아랫 칸의 수들의 집합이 일치하는지 보았다. 만약 일치하면 일치하는 집합을(윗 칸 수들의 집합이든 아랫 칸 수들의 집합이든 같으니까 아무거나 해도 됨) ans 리스트에 추가했다. 마지막으로는 ans리스트를 set으로 바꿔서, 중복을 제거한 후에 다시 리스트로 바꾸고 정렬을 했다.

import sys

input = sys.stdin.readline

N = int(input())

adj = [[] for _ in range(N + 1)]

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


# DFS로 탐색하다가 사이클 발생한 것들을 다 더하면 된다 !!

def dfs(num):
    if visited[num] == False:
        visited[num] = True
        for a in adj[num]:

            tmp_up.add(num)
            tmp_bottom.add(a)

            if tmp_up == tmp_bottom:
                ans.extend(list(tmp_bottom))
                return

            dfs(a)
    visited[num] = False


ans = []

for i in range(1, N + 1):
    visited = [False] * (N + 1)  # 위에 값 기준으로
    tmp_up = set()
    tmp_bottom = set()

    dfs(i)

ans = list(set(ans))
ans.sort()

print(len(ans))
for a in ans:
    print(a)