[백준] 1158: 요세푸스 문제 (python 파이썬)

1158번: 요세푸스 문제

해당 문제는 아래와 같이 단순히 구현만 한다면, 시간 초과가 발생한다.

n, k = map(int, input().split())

count = 1
index = 0
result = list()
check_list = [False for _ in range(n)]

while True:
    if not check_list[index]:
        if count == k:
            result.append(str(index+1))
            check_list[index] = True
            count = 1
        else:
            count += 1
    if index >= n-1:
        if False not in check_list:
            break
        index = 0
    else:
        index += 1

print('<'+', '.join(result)+'>')

• 해당 코드는 0~n-1번 원소에 False를 넣고, 빼낸 숫자들은 True로 바꿔준 코드다. 즉, 이미 나간 숫자들은 True, 아직 나가지 않은 숫자는 False 값을 가진다.
• 위 반복문의 종료 조건은 모든 리스트가 True 값을 가져야 한다는 것이다. 이 부분에서 많은 시간이 소요되어 시간 초과가 발생한다. N이 5,000으로 주어진다면, 5,000번의 탐색 이후(원돌기 이후), 다시 5,000개의 각각의 원소가 True인지 확인해야 한다. 이 경우 한 사이클만 25,000,000의 연산을 수행한다는 말이다.

그래서 생각한 방법이 이전에 풀었던 "1406: 에디터"문제의 풀이법이다. 내가 탐색하고 있는 인덱스 기준으로 왼쪽과 오른쪽을 다른 리스트에 보관하는 방식이다. 이전 문제의 풀이는 다음과 같다.

[백준] 1406: 에디터 (python 파이썬)

이 알고리즘을 이용해 큐를 사용하면 다음과 같이 코드를 작성할 수 있다.

from collections import deque

n, k = map(int, input().split())

left_queue = deque()
right_queue = deque([i+1 for i in range(n)])

count = 1
index = 0
result = list()

while True:
    if right_queue:
        if count == k:
            result.append(str(right_queue.popleft()))
            count = 1
        else:
            left_queue.append(right_queue.popleft())
            count += 1
    elif left_queue:
        right_queue = left_queue.copy()
        left_queue.clear()
    else:
        break

print('<'+', '.join(result)+'>')

하지만 해당 문제는 수학적으로 접근하면 더 쉽고 빠르게 풀 수 있다. 알고리즘은 다음과 같다.

• 1부터 n까지의 원소를 가진 리스트를 만든다.
• 탐색은 0번 인덱스에서 시작한다.
• 다음에 없어질 인덱스는 현재 탐색하고 있던 index에서 k-1을 더해주면 된다. 만약 k-1을 더해서 index가 사람 수를 넘어갔다면, index를 사람 수로 나눠 주고, 해당 결과의 나머지 값이 index가 된다. 예를 들어, 1~7번 사람이 주어졌고, k=3이라고 하자. 그럼 먼저 3번(2번 인덱스) 사람이 지워진다. 그 다음 6번(4번 인덱스, 3번 사람이 지워져 인덱스는 한 칸 당겨진다.)이 지워진다. 그리고  6번 인덱스가 지워져야 하는데, 현재 남은 인원은 5명이다. 따라서 그 2번(1번 인덱스) 사람을 지워야 한다.
• 위 과정을 반복하면 된다.

위 알고리즘을 바탕으로 코드를 작성하면 다음과 같다.

n, k = map(int, input().split())

index = 0
result = list()
check_list = [i+1 for i in range(n)]

while check_list:
    index += k-1
    num_person = len(check_list)
    if index >= num_person:
        index %= num_person
    result.append(str(check_list.pop(index)))
    
print('<'+', '.join(result)+'>')