[Python] 백준 9466번: 텀 프로젝트

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

9466번: 텀 프로젝트

 

import sys
sys.setrecursionlimit(300000)
input = sys.stdin.readline


def search(node):
    stack.append(node)
    visitInfo[node] = curSearchStart

    next = nums[node]
    nextVisitInfo = visitInfo[next]

    # 이번 탐색 때 방문한 경우
    if nextVisitInfo == curSearchStart:
        global notInCycle
        notInCycle -= len(stack) - stack.index(next)

    # 방문한 적이 없을 경우
    elif nextVisitInfo == 0:
        search(next)


if __name__ == '__main__':
    T = int(input())
    for _ in range(T):
        # 입력
        N = int(input())
        nums = [-1] + list(map(int, input().split()))

        # 탐색
        notInCycle = N
        visitInfo = [0]*(N+1)

        for i in range(1, N+1):
            # 방문한 적이 없을 경우
            if visitInfo[i] == 0:
                curSearchStart = i
                stack = []

                search(i)

        # 출력
        print(notInCycle)

 

시간 초과를 뚫기가 꽤 힘든 문제이다. DFS를 응용하여 풀었다.

 

나는 Boolean 타입의 isVisited을 쓰지 않았다. visitInfo에 방문했을 때의 탐색 시작 노드를 저장을 하였다. 가령, 2부터 탐색을 시작했고, 4에 방문하였다면, visitInfo[4] = 2이다. 이것을 처음에 0으로 초기화를 한다. 이렇게 함으로써, 방문한 적이 있는지에 대한 정보(0 또는 1 이상)와 어떤 탐색 때 방문하였는지를 모두 알 수 있다.

 

탐색을 시작하면, 먼저 stack에 현재 노드를 넣는다. stack에는, 현재 탐색 때 방문했던 노드들을 기록한다. 그리고 visitInfo에 현재 탐색 시작점(curSearchStart)를 넣어준다.

다음으로 visitInfo를 통해, 다음 노드가 이번 탐색 때 방문했는지를 알 수 있는데, 이 경우 이 문제의 출력값인 notInCycle(N으로 초기화된다.)에 적절한 값을 넣어준다. 가령 방문을 1->2->3->4->5->2 순서로 했다고 하자. 마지막 2에서 해당 조건문에 걸리게 되는데, 이 때 next는 2이고, stack은 [1, 2, 3, 4, 5]이다. 여기서 cycle은 2, 3, 4, 5에 해당한다. 따라서, (스택의 길이) - (next의 인덱스)를 빼준다.

그리고 방문한 적이 없을 경우에는 DFS 탐색을 해준다.

 

포인트는 visitInfo에 curSearchStart를 넣음으로써 어떤 탐색 때 방문하였는지, 그리고 동시에 방문했는지를 알 수 있다는 점과, stack에 방문 기록을 저장하여, 현재 탐색 때 방문한 적이 있는 경우 싸이클의 길이를 알 수 있다는 점이다.