[PS] BOJ 9466 / 텀 프로젝트

[PS] BOJ 9466 / 텀 프로젝트
문제 링크: https://www.acmicpc.net/problem/9466
Thumbnail: Photo by Vitaly Gariev (Unsplash)

풀이

학생들이 팀을 이룰 수 있는 조건은 다음과 같습니다.

학생들이(s1, s2, ..., sr)이라 할 때, r=1이고 s1이 s1을 선택하는 경우나, s1이 s2를 선택하고, s2가 s3를 선택하고,..., sr-1이 sr을 선택하고, sr이 s1을 선택하는 경우에만 한 팀이 될 수 있다.

다시 말하면, 사이클이 이루어지는 경우에만 팀이 됩니다. DFS로 각 학생들을 탐색하면서, 사이클을 찾아 길이를 계산해주면 됩니다.

DFS

기본적으로, 학생들이 선호하는 팀을 탐색하는 방식은 DFS입니다. 학생들이 서로 가리키는 순으로 먼저 탐색한 뒤, 다른 학생을 탐색하기 때문입니다.

먼저 학생들이 서로 선호하는 학생에 대한 정보를 입력으로 받습니다.

이후, 1번 학생부터 N번 학생까지 DFS를 진행해줍니다. 이때, 한 번 방문했던 학생 정보는 더 이상 방문할 필요가 없으니 visited 배열을 통해 검사합니다.

input = open(0).readline
preferred_teams = [0] * 100_001           # 입력 데이터를 저장할 배열
visited = [False for _ in range(100_001)] # 방문 여부를 저장할 배열
teamed_cnt = 0 # 팀을 맺은 학생들의 수를 기록할 변수. dfs 내부에서 전역변수로 접근한다.

def solution():
    global teamed_cnt
    N = int(input())
    teamed_cnt = 0
    for idx, value in enumerate(map(int, input().split()), 1):
        preferred_teams[idx] = value
        visited[idx] = False
    
    path = [0 for _ in range(N + 1)] # DFS를 진행하며 방문한 정점의 순서를 기록할 배열
    for i in range(1, N + 1):
        if not visited[i]:
            visited[i] = True
            path[0] = i
            dfs(i, 0, path)
    
    print(N - teamed_cnt)

for _ in range(int(input())):
    solution()

입력과 DFS

사이클 찾기

먼저, 기본적인 DFS를 구현합니다.

def dfs(n, depth, path):
    next_n = preferred_teams[n]
    if not visited[next_n]:
        visited[next_n] = True
        path[depth + 1] = next_n
        dfs(next_n, depth + 1, path)
        path[depth + 1] = 0

DFS 구현

이제, 사이클을 찾아야합니다. 사이클의 경우 기존에 방문했던 노드를 다시 방문할 때 발생하므로, visited[next_n]이 참인 경우에 사이클을 찾으면 됩니다.

한 번 방문했던 정점에 다시 방문 시, path 배열을 순회하며 다음에 방문하려 했던 정점인 next_n과 일치하는 정점을 찾습니다. 이후, 발견한 사이클의 길이만큼 팀을 맺은 학생의 수를 더해줍니다.

def dfs(n, depth, path):
    next_n = preferred_teams[n]
    ...
    else:
        for prev in range(depth):
            if path[prev] == next_n:
                global teamed_cnt
                teamed_cnt += depth + 1 - prev
                return

사이클 찾기

모든 학생에 대해 DFS를 완료한 뒤, 전체 학생 수에서 팀을 맺은 학생의 수를 빼면 됩니다.

약간의 트릭

입력 데이터를 받을 때, 자기 자신을 팀으로 선호하는 학생의 경우 바로 사이클을 찾을 수 있으므로 즉시 처리해주면 이후 DFS의 호출을 조금이나마 줄일 수 있습니다.

def solution():
    global teamed_cnt
    N = int(input())
    teamed_cnt = 0
    for idx, value in enumerate(map(int, input().split()), 1):
        preferred_teams[idx] = value
        visited[idx] = False
        if idx == value: # 자기 자신을 팀으로 선호하는 경우는 즉시 사이클을 판단할 수 있으므로 사전에 처리해 둔다.
            visited[idx] = True
            teamed_cnt += 1
    ...

전체 코드

from sys import setrecursionlimit
setrecursionlimit(10**5)
input = open(0).readline
preferred_teams = [0] * 100_001 # 전역 변수로 선언
visited = [False for _ in range(100_001)]
teamed_cnt = 0

def dfs(n, depth, path):
    next_n = preferred_teams[n]
    if not visited[next_n]:
        visited[next_n] = True
        path[depth + 1] = next_n
        dfs(next_n, depth + 1, path)
        path[depth + 1] = 0
    else:
        for prev in range(depth):
            if path[prev] == next_n:
                global teamed_cnt
                teamed_cnt += depth + 1 - prev
                return

def solution():
    global teamed_cnt
    N = int(input())
    teamed_cnt = 0
    for idx, value in enumerate(map(int, input().split()), 1):
        preferred_teams[idx] = value
        visited[idx] = False
        if idx == value: # 자기 자신을 팀으로 선호하는 경우는 즉시 사이클을 판단할 수 있으므로 사전에 처리해 둔다.
            visited[idx] = True
            teamed_cnt += 1
    
    path = [0 for _ in range(N + 1)]
    for i in range(1, N + 1):
        if not visited[i]:
            visited[i] = True
            path[0] = i
            dfs(i, 0, path)
    
    print(N - teamed_cnt)

for _ in range(int(input())):
    solution()

solution.py