[백준 2668] 숫자고르기

문제

 

www.acmicpc.net/problem/2668

2668번: 숫자고르기

 

처음 작성한 코드

import sys

N = int(sys.stdin.readline())
arr = {i:int(sys.stdin.readline()) for i in range(1, N+1)} 
visited = [False] * (N+1)
result = []

for i in range(1, N+1):
    tmp = []
    if not visited[i]:
        visited[i] = True
        tmp.append(i)
        while tmp:
            if not visited[arr[tmp[-1]]]:
                visited[arr[tmp[-1]]] = True
                tmp.append(arr[tmp[-1]])
            else:
                if tmp[0] == arr[tmp[-1]]:
                    for j in tmp:
                        result.append(j)
                else:
                    for j in tmp:
                        visited[j] = False
                tmp = []

print(len(result))
for i in sorted(result):
    print(str(i))

 

예제 입력을 기준으로 설명하겠다.

1. 1->3, 3->1 이면 루프가 생기고 이런 루프가 생기면 result에 넣는 방식으로 작성. 5->5도 루프가 생기니까 result에 넣음
2. while문 안에 첫 if-else 문은 현재 번호와 연결되어있는 번호가 방문이 되어있는지를 체크한다. 만약 방문이 안됐으면 방문됐다고 체크하고 tmp에 연결되어있는 번호를 넣는다. 이렇게 계속 하다보면 연결되어있는 모든 노드를 계속 들어가면서 방문하게 된다.
3. else(연결되어있는 번호가 이미 방문됐다면)를 만나서 만약 시작점과 지금 가려는 번호(방문되었다고 체크되서 else로 빠진 번호)가 같으면 루프를 생성하니까 result에 tmp에 있는 번호를 다 넣는다. ex) 1->3->1(arr[tmp[-1]])일때 tmp[0](시작점)과 arr[tmp[-1]](지금 가려던 번호)가 같으니 tmp에 있는 1과 3을  result에 넣는다.
4. 만약 시작번호와 지금 번호가 같지 않다면 방문되어있다는 걸 다시 false로 바꿔줘야 그 뒤 진행을 다시 할 수 있다. ex) 4->5->5 일때 4와 5를 visited false로 해줘야 5->5일때를 잡을 수 있다 왜냐면 제일 큰 for 문에서 not visited상대로만 진행하기 때문에..

코드 리뷰 후

import sys

def DFS(i, start):
    visited[i] = True

    if not visited[arr[i]]:
        DFS(arr[i],start)
    elif visited[arr[i]] and arr[i] == start:
        result.append(i)


N = int(sys.stdin.readline())
arr = {i:int(sys.stdin.readline()) for i in range(1, N+1)}
result = []

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

print(len(result))
for i in sorted(result):
    print(i)

처음 작성한 코드를 재귀로 푼 모습이다!

 

깨달은 점

• 마지막에 출력할 때 for문 돌면서 프린트하는 걸 까먹어서 '\n'.join(result)이런식으로 했는데 결과는 똑같던데 백준에서는 틀렸다고 해서 그거때문에 시간 날렸음!!!
• 생각하는데 오래걸리지 않았는데 구현하는데 오래걸렸다 왤까?