[BOJ/Python] 1158 - 요세푸스 문제

출력이 조금 까다로웠던 문제다. 

1. 시간 초과

from sys import stdin
n,k = map(int,stdin.readline().split())
arr = [x for x in range(1,n+1)]
stack = [x for x in range(1,n+1)]
deleted = []
for i in range(1,n+1):
    idx = k*i-1
    if idx>len(stack)-1:
        for i in arr:
            if i in deleted:
                arr.remove(i)
        while True:
            if len(stack)>idx:
                break
            stack.extend(arr)
    if stack[idx] not in deleted:
        deleted.append(stack[idx])
print('<',', '.join([str(x) for x in deleted]),'>',sep='')

처음에는

1 2 3 4 5 6 7  - 1 2 4 5 7 + 1 4 5 +  1 4 +  1 4 4 4 4  이런 식으로, 지워야 될 k번째의 수를 지워 가면서 리스트에 append하는 방식을 생각했지만, delete와 extend를 위해 반복문이 두 번 쓰여야 하다 보니 예상대로 시간이 초과되었다.

2. 통과

from sys import stdin
n, k = map(int, stdin.readline().split())
arr = [x for x in range(1, n+1)]
idx = k-1 # k번째의 수를 인덱스로 보기 위해 1을 빼주고
deleted = []
while arr:
    while idx > len(arr)-1: 
        idx -= len(arr) # 만약 남은 리스트의 전체 길이보다 k-1 = idx이 크다면 
        				#idx의 값에서 현재 arr의 길이만큼 뺀다.
    deleted.append(arr.pop(idx))
    idx+=(k-1)
print('<', ', '.join([str(x) for x in deleted]), '>', sep='')

요세푸스 문제

7 3을 입력했을 때 어떻게 진행되는지 출력한 결과이다.

 

[1,2,4,5,7] 리스트에서 인덱스 6을 슬라이싱한다고 생각하면 간단하다.

리스트의 길이-1 보다 인덱스가 작거나 같아질 때까지 해당 arr의 전체 길이만큼 인덱스에서 빼면 된다.