Lưu ý: các đoạn code trong bài viết này sử dụng 0-indexing và khoảng nửa mở $[l, r)$.
Ta có thể coi $a$ như một tập hợp các đoạn đã được sắp xếp.
Ví dụ với mảng $a = [1, 6, 3, 4, 2, 5, 7]$:
- Ban đầu: $\mathcal A = [\\{1\\}, \\{6\\}, \\{3\\}, \\{4\\}, \\{2\\}, \\{5\\}, \\{7\\}]$
- Sau truy vấn `0 3 6`: $\mathcal A = [\\{1\\}, \\{6\\}, \\{2, 3, 4, 5\\}, \\{7\\}]$
- Sau truy vấn `1 5 7`: $\mathcal A = [\\{1\\}, \\{6\\}, \\{2, 3, 4\\}, \\{7, 5\\}]$
- Sau truy vấn `0 1 4`: $\mathcal A = [\\{1, 2, 3, 6\\}, \\{4\\}, \\{7, 5\\}]$
Khi đó bài toán này quy về thành:
- Gộp hai dãy $S_1, S_2$ đã được sắp xếp thành dãy $S=S_1\cup S_2$
- Tách dãy $S$ thành một dãy $S_1$ gồm $k$ phần tử nhỏ nhất của $S$ và dãy $S_2$
gồm các phần tử còn lại
Để biểu diễn dãy các số nguyên không âm, ta sẽ sử dụng cấu trúc cây phân đoạn thưa.
```cpp
struct node {
node *l = 0, *r = 0;
// khi khởi tạo dãy đơn, gán a[i] vào đây
int val = 0;
int cnt = 0;
void recalc() {
// khi thực hiện truy vấn tại vị trí `pos`,
// ta sẽ tách đoạn chứa `pos` để lấy riêng đoạn [pos, pos+1)
// Sau khi tách, gốc của đoạn này là một dãy đơn,
// do đó `val` tại gốc chính là giá trị của phần tử duy nhất đó
val = (l ? l->val : 0) + (r ? r->val : 0);
cnt = (l ? l->cnt : 0) + (r ? r->cnt : 0);
}
};
```
## Gộp hai cây phân đoạn
Để gộp hai cây phân đoạn $\mathcal T_1, \mathcal T_2$ thành $\mathcal T$, ta chỉ việc đệ quy gộp
các nút con của chúng. Tính đúng đắn được đảm bảo vì tại mỗi nút, giá trị $\text{cnt}$ luôn bằng
tổng số phần tử trong hai cây con; do đó khi gộp, ta chỉ cần cộng dồn số lượng từ hai cây con
tương ứng. Như vậy, $\mathcal T = \mathcal T_1 \cup \mathcal T_2$.
Sau đây là hàm để gộp hai cây phân đoạn:
```cpp
merge(t1, t2) {
if (!t1 || !t2) return t1 ? t1 : t2;
t1->l = merge(t1->l, t2->l);
t1->r = merge(t1->r, t2->r);
t1->recalc();
return t1;
}
```
Mỗi lần gọi $\text{merge}$ ta sẽ tiêu thụ các nút của $\mathcal T_2$. Sau khi một nút của $\mathcal T_2$,
được gộp vào $\mathcal T_1$ nó sẽ không được truy cập lại trong những thao tác sau. Do đó, mặc dù
một lần gộp riêng lẻ có thể tốn đến $\mathcal O(n)$ thao tác, tổng chi phí toàn bộ quá trình sẽ
là $\mathcal O(\sum|\mathcal T|) = \mathcal O(n\log U)$* ($U$ là giá trị tối đa mà cây phân đoạn
chứa được. Ở đây, $U = n$). Vì vậy $\text{merge}$ chỉ tốn trung bình $\mathcal O(\log U)$ cho mỗi lần gọi.
*: Đúng hơn là $\mathcal O((n+q)\log U)$, giải thích ở phần tiếp theo.
## Tách cây phân đoạn
Ta có thể chia $\mathcal T$ thành hai cây con: $\mathcal T_1$ gồm $k$ phần tử nhỏ nhất của
$\mathcal T$, và $\mathcal T_2$ gồm các phần tử còn lại, bằng cách đi xuống cây theo giá trị
của $\text{cnt}$.
```cpp
split(t, k) {
if (k == 0) return {0, t};
if (k == t->cnt) return {t, 0};
szl = (t->l ? t->l->cnt : 0);
node *t2 = new node;
if (k < szl) {
(t2->l, t->l) = split(t->l, k);
swap(t, t2); // t2 đang chứa k phần tử nhỏ nhất
} else {
(t->r, t2->r) = split(t->r, k - szl);
}
t->recalc(), t2->recalc();
return {t, t2};
}
```
Với mỗi lần gọi $\text{split}$ sẽ tạo $\mathcal O(\log U)$ nút khiến tổng chi phí trở thành
$\mathcal O((n+q)\log U)$.
## Xử lý các truy vấn sắp xếp
Ta sẽ lưu:
- `vector<node*> roots(n)` chứa gốc của các cây phân đoạn
- `vector<bool> rev(n)` để biết đoạn được sắp xếp tăng dần (`false`) hay giảm dần (`true`)
- `std::set avail` để quản lý những vị trí có tồn tại gốc (để tiện hơn nên có $n$ trong
$\text{avail}$)
Giờ ta chỉ việc tách các dãy tại hai đầu $l, r$ và gộp tất cả các dãy trong $[l, r)$. Cần chú ý
khi tách các dãy sắp xếp giảm dần.
```cpp
cut(pos) {
l = *--avail.upper_bound(pos);
if (l == pos || l >= n) return;
if (!rev[l]) {
(roots[l], roots[pos]) = split(roots[l], pos - l);
} else {
(roots[pos], roots[l]) = split(roots[l], roots[l]->cnt - (pos - l));
}
avail.insert(pos);
}
sort(descending, l, r) {
cut(l), cut(r);
for pos: l < pos < r {
roots[l] = merge(roots[l], roots[pos]);
avail.erase(pos);
}
rev[l] = descending;
}
query(pos) {
cut(p), cut(p + 1);
return roots[p]->val;
}
```
Một lưu ý là ta đang làm việc trên tập hợp nhưng trong bài, $a$ có thể có các giá trị trùng nên
ta cần rời rạc hóa mảng $a$.
Tổng độ phức tạp thời gian: $\mathcal O((n + q)\log U)$