Chương trình thiết kế hàng đợi di chuyển phần tử gần đây nhất về cuối trong Python

Mô tả bài toán

Thiết kế hàng đợi khởi tạo với số 1 đến n. Hàm solve(k) di chuyển phần tử tại vị trí k (1-indexed) về cuối hàng đợi và trả về giá trị hiện tại ở cuối hàng đợi sau khi di chuyển.

Ví dụ

n = 5, các lần gọi solve(5), solve(2), solve(3), solve(1) cho kết quả 5, 2, 4, 1.

Giải thuật (Sqrt Decomposition)

1. Chia mảng thành các block kích thước ~sqrt(n).
2. Dùng bisect_right tìm block chứa vị trí k.
3. Pop phần tử, cập nhật index các block sau, thêm phần tử vào block cuối.
4. Nếu block cuối quá lớn, tách block mới.
5. Nếu block cũ rỗng, xóa block đó.

Triển khai Python

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))

Đầu vào

queue = TestQueue(5); print(queue.solve(5)); print(queue.solve(2)); print(queue.solve(3)); print(queue.solve(1))

Đầu ra

5 2 4 1