문제 소개
정수 1부터 n까지의 값으로 초기화되는 특별한 큐(queue)를 설계해야 한다고 가정해 보겠습니다. 이 큐는 함수가 호출될 때마다 입력으로 전달된 위치에 있는 요소를 큐의 맨 뒤로 이동시키고, 이동 작업이 끝난 후 현재 큐의 맨 뒤에 있는 값을 반환하는 기능을 수행합니다.
예를 들어 n = 5로 큐를 초기화하면 [1, 2, 3, 4, 5] 상태가 되며, 이동 명령이 위치 5, 2, 3, 1 순서로 호출될 때의 동작은 다음과 같습니다.
- solve(5): 5번째 요소(값 5)를 맨 뒤로 이동 → [1, 2, 3, 4, 5], 마지막 값 5 반환
- solve(2): 2번째 요소(값 2)를 맨 뒤로 이동 → [1, 3, 4, 5, 2], 마지막 값 2 반환
- solve(3): 3번째 요소(값 4)를 맨 뒤로 이동 → [1, 3, 5, 2, 4], 마지막 값 4 반환
- solve(1): 1번째 요소(값 1)를 맨 뒤로 이동 → [3, 5, 2, 4, 1], 마지막 값 1 반환
따라서 최종 출력은 5, 2, 4, 1이 됩니다.
접근 방법: √n 분할(Sqrt Decomposition)
파이썬 리스트의 pop()과 append()만으로 단순 구현하면 한 번의 연산에 최대 O(n)의 시간이 소요됩니다. 데이터 크기가 커지면 비효율적이므로, 여기서는 √n 분할 기법을 활용해 성능을 개선합니다.
핵심 아이디어는 전체 큐를 크기 약 √n짜리 여러 블록으로 나누고, 각 블록의 첫 번째 값을 별도의 인덱스 배열(index)에 기록해 두는 것입니다. 이렇게 하면 요소를 찾고 이동할 때 해당 블록 내부에서만 연산을 수행하면 되므로, 각 명령을 평균 O(√n)에 처리할 수 있습니다.
알고리즘 단계
- i := bisect_right(index, k) − 1 → 값 k가 속한 블록의 위치를 찾습니다.
- x := data[i]에서 (k − index[i])번째 요소를 꺼냅니다(pop).
- i+1부터 index 배열의 끝까지 모든 index[ii] 값을 1씩 감소시킵니다.
- 마지막 블록의 크기가 nn(≈√n) 이상이면 새 블록을 만들고 index에 n을 추가합니다.
- 꺼낸 값 x를 마지막 블록(data[-1])의 끝에 삽입합니다.
- 블록 data[i]가 비어 있으면 해당 블록과 index 항목을 제거합니다.
- x를 반환합니다.
파이썬 구현 예제
from bisect import bisect_right
from math import sqrt
class TestQueue:
def __init__(self, n):
self.n = n
self.nn = int(sqrt(n))
self.data = []
self.index = []
for i in range(1, n + 1):
ii = (i - 1) // self.nn
if ii == len(self.data):
self.data.append([])
self.index.append(i)
self.data[-1].append(i)
def solve(self, k):
i = bisect_right(self.index, k) - 1
x = self.data[i].pop(k - self.index[i])
for ii in range(i + 1, len(self.index)):
self.index[ii] -= 1
if len(self.data[-1]) >= self.nn:
self.data.append([])
self.index.append(self.n)
self.data[-1].append(x)
if not self.data[i]:
self.data.pop(i)
self.index.pop(i)
return x
queue = TestQueue(5)
print(queue.solve(5))
print(queue.solve(2))
print(queue.solve(3))
print(queue.solve(1))
입력
queue = TestQueue(5) print(queue.solve(5)) print(queue.solve(2)) print(queue.solve(3)) print(queue.solve(1))
출력
5 2 4 1
코드 동작 원리
__init__ 메서드: 블록 크기 nn을 int(sqrt(n))로 계산한 뒤, 1부터 n까지의 값을 순회하며 각 값이 속할 블록 번호 ii = (i−1)//nn을 구합니다. 새 블록이 필요해지면 data에는 빈 리스트를, index에는 해당 블록의 시작 값을 추가하고, 값을 마지막 블록에 차례대로 채워 넣습니다.
solve 메서드: bisect_right를 이용해 k가 속한 블록 i를 찾고, k − index[i]로 블록 내 상대 위치를 계산해 요소를 꺼냅니다. 요소가 하나 제거되었으므로 이후 블록들의 시작 인덱스를 1씩 줄이고, 마지막 블록이 가득 차면 새 블록을 생성한 뒤 꺼낸 값 x를 맨 뒤에 붙입니다. 마지막으로 원래 블록이 비었다면 불필요한 블록을 정리하고 x를 반환합니다.
이 구조 덕분에 요소의 검색·삭제·삽입이 대부분 하나의 블록 안에서 일어나며, 전체 연산 복잡도는 O(√n)으로 단순 리스트 구현(O(n))보다 훨씬 효율적입니다.