문제

이번 가을학기에 ‘문제 해결’ 강의를 신청한 학생들은 텀 프로젝트를 수행해야 한다. 프로젝트 팀원 수에는 제한이 없다. 심지어 모든 학생들이 동일한 팀의 팀원인 경우와 같이 한 팀만 있을 수도 있다. 프로젝트 팀을 구성하기 위해, 모든 학생들은 프로젝트를 함께하고 싶은 학생을 선택해야 한다. (단, 단 한 명만 선택할 수 있다.) 혼자 하고 싶어하는 학생은 자기 자신을 선택하는 것도 가능하다.

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

예를 들어, 한 반에 7명의 학생이 있다고 하자. 학생들을 1번부터 7번으로 표현할 때, 선택의 결과는 다음과 같다.

1 2 3 4 5 6 7
3 1 3 7 3 4 6

위의 결과를 통해 (3)과 (4, 7, 6)이 팀을 이룰 수 있다. 1, 2, 5는 어느 팀에도 속하지 않는다.

주어진 선택의 결과를 보고 어느 프로젝트 팀에도 속하지 않는 학생들의 수를 계산하는 프로그램을 작성하라.

www.acmicpc.net/problem/9466

  • 입력

첫째 줄에 테스트 케이스의 개수 T가 주어진다. 각 테스트 케이스의 첫 줄에는 학생의 수가 정수 n (2 ≤ n ≤ 100,000)으로 주어진다. 각 테스트 케이스의 둘째 줄에는 선택된 학생들의 번호가 주어진다. (모든 학생들은 1부터 n까지 번호가 부여된다.)

  • 출력

각 테스트 케이스마다 한 줄에 출력하고, 각 줄에는 프로젝트 팀에 속하지 못한 학생들의 수를 나타내면 된다.

과정

그래프의 사이클을 구하는 문제다. 사이클을 구하는 문제는 처음이라서 조금 어려움이 있었다. 팀에 속하지 못한 학생들을 저장하는 student 집합을 만들어 문제를 해결하려 했지만 시간 초과가 떴다.

import sys

t = int(sys.stdin.readline())
while t:
    n = int(sys.stdin.readline())
    s = [0] + list(map(int, sys.stdin.readline().split()))
    # student: 팀에 속하지 못한 학생들을 저장
    student = set()

    for i in range(1, n+1):
        tmp = []
        j = i
        while True:
            # j가 tmp에 이미 존재하는 경우 사이클을 찾은 것
            if j in tmp:
                for k in range(0, n):
                    if tmp[k] == j:
                        break
                    # 사이클에 속하지 않는 학생들을 집합에 추가
                    else:
                        student.add(tmp[k])
                break
            else:
                tmp.append(j)
                j = s[j]

    print(len(student))
    t -= 1

최종 코드

수정한 코드는 다음과 같다. 그룹에 속하지 못한 학생들을 집합에 추가하는 대신 group이라는 배열을 만들어 학생들에게 그룹 번호를 부여한다.

해당 그룹에서 사이클을 찾으면, 사이클을 돌면서 그룹을 이룬 학생들에게는 -1의 번호를 부여한다.

마지막에 번호가 양수인 학생의 수를 구하면 답을 구할 수 있다.

import sys

t = int(sys.stdin.readline())
while t:
    n = int(sys.stdin.readline())
    s = [0] + list(map(int, sys.stdin.readline().split()))
    group = [0]*(n+1)

    g_num = 1
    for i in range(1, n+1):
        if group[i] == 0:
            j = i
            # 사이클을 찾을 때까지 그룹 번호를 부여
            while group[j] == 0:
                group[j] = g_num
                j = s[j]
            # 사이클을 찾으면 사이클에 속하는 학생들에게 그룹 번호 -1을 부여
            while group[j] == g_num:
                group[j] = -1
                j = s[j]
            g_num += 1

    cnt = 0
    for i in range(1, n+1):
        if group[i] > 0 :
            cnt += 1

    print(cnt)
    t -= 1

Enhanced

DFS로 푸는 방법도 있다. 일단 recursion limit을 설정해야 한다.

앞선 코드에서 그룹 번호를 부여한 것처럼 이번엔 팀을 짜는 것이다. 만약 사이클에 속한다면 ans 배열에 팀에 속한 학생들을 추가한다.

마지막에 전체 학생 숫자에서 팀에 속한 학생들의 숫자(len(ans))를 빼주면 답을 구할 수 있다.

import sys
sys.setrecursionlimit(1000000)

def dfs(start):
    global team, ans, visit
    team.append(start)
    visit[start] = 1
    next = s[start]

    # 다음 노드가 방문하지 않은 노드인 경우 dfs 재귀 호출
    if not visit[next]:
        dfs(next)
    # 사이클을 찾았다면
    else:
        if next in team:
            # 사이클에 속한 학생들을 ans 배열에 추가
            ans += team[team.index(next):]
        return

t = int(sys.stdin.readline())
while t:
    n = int(sys.stdin.readline())
    s = [0] + list(map(int, sys.stdin.readline().split()))
    visit = [0]*(n+1)
    ans = []

    for i in range(1, n+1):
        if not visit[i]:
            team = []
            dfs(i)

    print(n-len(ans))
    t -= 1

참고 사이트