Python에서 재귀 인덱싱으로 요소 집합의 개수를 계산하는 프로그램

문제 소개

숫자로 이루어진 리스트 A와 정수 k가 주어졌을 때, 다음과 같은 형태의 새로운 집합을 만들어야 합니다.

{A[k], A[A[k]], A[A[A[k]]], ...}

이 과정은 인덱스가 리스트 범위를 벗어나기 직전까지 반복됩니다. 최종적으로 이 집합의 크기를 구하는 것이 목표이며, 만약 탐색 도중 사이클(cycle)이 발견된다면 -1을 반환해야 합니다.

예시

입력이 A = [1,2,3,4,5,6,7], k = 1이라고 가정해 보겠습니다.

• A[1] = 2
• A[2] = 3
• A[3] = 4
• A[4] = 5
• A[5] = 6
• A[6] = 7

따라서 생성되는 집합은 {2, 3, 4, 5, 6, 7}이며, 집합의 크기는 6입니다.

해결 접근 방법

이 문제는 다음 단계를 통해 해결할 수 있습니다.

• 방문한 값을 저장할 빈 집합 seen을 생성합니다.
• k가 리스트 A의 길이보다 작은 동안 반복합니다.
  • 만약 A[k]가 이미 seen에 존재한다면, 사이클이 발생한 것이므로 -1을 반환합니다.
  • 그렇지 않으면 A[k]를 seen에 추가합니다.
  • k의 값을 A[k]로 갱신하여 다음 위치로 이동합니다.
• 반복이 종료되면 seen의 크기를 반환합니다.

사이클 여부를 판별하는 핵심은 이미 방문한 값들을 집합(set)에 기록해 두는 것입니다. 집합은 중복을 허용하지 않으므로, 동일한 값이 다시 나타나는 순간 무한 루프에 빠진다고 판단할 수 있습니다.

구현 코드

class Solution:
    def solve(self, A, k):
        seen = set()
        while k < len(A):
            if A[k] in seen:
                return -1
            seen.add(A[k])
            k = A[k]
        return len(seen)

ob = Solution()
print(ob.solve([1,2,3,4,5,6,7], 1))

입력

[1,2,3,4,5,6,7], 1

출력

6