2668. 숫자고르기 (Python)

2668번: 숫자고르기

 

풀이

 dfs를 통해서 사이클이 존재하는지를 체크해주어 해결할 수 있는 문제다.

 

 먼저 입력을 받으면서 뽑으면 안되는 정수를 걸러준다. 위의 예를 보면 아래 숫자 목록에서 2와 7은 존재하지 않는다. 그렇기 때문에 위의 숫자 목록에서 2와 7은 굳이 dfs를 하지 않아도 일치하는 집합을 만들 수 없음을 알 수 있다.

 

 그 후 차례대로 숫자 1부터 dfs 탐색을 진행한다. 1부터 숫자를 뽑아가며 배열에 저장하고 사이클이 이루어지는지를 확인한다. 예를 들어 숫자 1을 뽑으면 진행 순서는 다음과 같다. 1 -> 3 -> 1  이렇게 되면 숫자 1에서 출발하여 사이클이 이루어지고 있으므로 1과 3을 뽑을 수 있는 숫자로 체크한다.

 다음으로 2는 처음에 뽑을 수 없는 숫자로 체크하였고 3은 이미 1에서 dfs를 통해서 사이클이 이루어지는 숫자임을 확인하였기 때문에 dfs를 진행하지 않고 숫자 4로 넘어간다. 숫자 4를 뽑으면 4 -> 5 -> 5와 같이 진행된다. 이렇게 되면 숫자 5에서 출발하여 사이클이 이루어지고 있으므로 숫자 5를 뽑으면 일치하는 집합을 만들 수 있음을 알 수 있다.

 

 위와 같은 방식으로 처음부터 마지막까지 dfs를 진행하면 최대로 많이 뽑는 개수를 알 수 있다.

 

코드

import sys
import heapq
input = sys.stdin.readline


# dfs 함수 
# numCycle 배열 : 현재까지 뽑은 숫자를 기록해 놓는 배열 
# numIdx 배열 : 현재 숫자 n에 대하여 numCycle 배열의 몇 번째에 있는지를 기록해 놓는 배열 
#              초기값은 -1로 되어 있으며 만약 numIdx[n]의 값이 -1이 아니라면 정수 n이 numCycle에 이미 존재한다는 뜻이며
#              이는 numCycle 배열에 cycle이 생겼음을 의미한다.
def dfs(n, idx):
    global res

    # numIdx[n]이 -1이 아니라면 사이클이 생겼음을 의미
    # numIdx[n]부터 numCycle의 마지막 숫자까지가 사이클에 해당한다.
    if numIdx[n] != -1:
        for i in range(len(numCycle)):
            if numIdx[n] <= i:
                res += 1
                heapq.heappush(resQ, numCycle[i])
    
            canPick[numCycle[i]] = False # 이후에 dfs를 돌지 않도록 각 정수에 대하여 False로 설정

        return

    next = nums[1][n] # 다음으로 탐색할 숫자

    if not canPick[next]: # dfs를 이미 수행한 숫자라면 곧바로 종료한다.
        return

    numCycle.append(n) # numCycle 배열에 현재 숫자 n을 추가한다.
    numIdx[n] = idx # n이 numCycle의 idx번째에 있음을 numIdx[n]에 기록한다.

    dfs(nums[1][n], idx+1)


if __name__ == '__main__':
    N = int(input())
    nums = [[i for i in range(N+1)] for _ in range(2)]
    canPick = [False for _ in range(N+1)] # dfs를 이미 수행했거나 수행할 필요가 없음을 체크하기 위한 배열 

    for i in range(1, N+1):
        tmp = int(input())
        nums[1][i] = tmp
        canPick[tmp] = True # 입력 과정에서 True로 설정되지 않은 숫자는 dfs를 수행할 필요가 없다

    res = 0
    resQ = []
    for i in range(1, N+1):
        if not canPick[i]: # dfs를 수행했거나 수행할 필요가 없는 숫자는 건너뛴다.
            continue

        numIdx = [-1 for _ in range(N+1)] # 초기값은 -1로 설정
        numCycle = []
        dfs(i, 0)

    print(res)
    while resQ:
        print(heapq.heappop(resQ))