[백준] 1158번 요세푸스 문제

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

1158번: 요세푸스 문제

문제

요세푸스 문제는 다음과 같다.

1번부터 N번까지 N명의 사람이 원을 이루면서 앉아있고, 양의 정수 K(≤ N)가 주어진다. 이제 순서대로 K번째 사람을 제거한다. 한 사람이 제거되면 남은 사람들로 이루어진 원을 따라 이 과정을 계속해 나간다. 이 과정은 N명의 사람이 모두 제거될 때까지 계속된다. 원에서 사람들이 제거되는 순서를 (N, K)-요세푸스 순열이라고 한다. 예를 들어 (7, 3)-요세푸스 순열은 <3, 6, 2, 7, 5, 1, 4>이다.

N과 K가 주어지면 (N, K)-요세푸스 순열을 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 N과 K가 빈 칸을 사이에 두고 순서대로 주어진다. (1 ≤ K ≤ N ≤ 5,000)

 

출력

예제와 같이 요세푸스 순열을 출력한다.

 

 

풀이

 

for i in range(1, N+1):
    q.append(i)

while q:
    if count == M-1:
        answer.append(q[idx])
        q.pop(idx)
        if idx == len(q):	#제거한 idx가 마지막 번호인 경우에도 0으로 초기화
            idx = 0
        count = 0


    count += 1			#몇 칸 떨어져있는지 체크
    idx += 1			#제거할 index

    if idx == len(q):	#끝까지 간다면 0으로 초기화
        idx = 0

하나씩 체크하면서 남은 배열의 길이까지 가면 처음으로 되돌아오는 방식으로 구현했더니 시간초과가 발생하였다.
입력이 5000개이기 때문에 5000번을 돌아도 2500만으로 시간 안에 들어올 것이라고 생각했는데 아마 list를 이용해서 배열이 앞으로 당겨지는 연산을 최대 N번 한다고 가정하면 1250억이 되어 시간 안에 안되는 것 같다.

 

전체 코드

import sys

N, M = map(int, sys.stdin.readline().split())
q = []
answer = []
numbers = [i for i in range(1, N+1)]
idx = 0

for i in range(N):
    idx += M-1
    if idx >= len(numbers):
        idx = idx % len(numbers)

    answer.append(numbers.pop(idx))
print('<' + ', '.join(map(str, answer)) + '>')

그래서 N개의 배열에 제거를 수행해서 당겨지는 연산이 N번 수행되더라도 풀 수 있게 idx를 움직이는 연산을 1로 수정하였다. idx에 이동할 만큼 바로 더한 뒤 남은 list의 길이보다 큰 경우 modular연산을 이용하여 나머지만 얻는 것으로 0으로 다시 온 뒤 체크하는 것과 동일하게 작동하게 하였다.