Chương trình C++ tìm phần tử lớn thứ k (Order Statistic Tree)

Sử dụng cây BST bổ sung trường rank (thứ tự) để tìm phần tử lớn thứ k. Rank được gán bằng duyệt in-order (trái-gốc-phải).

Thuật toán

Hàm Insert(root, d): Chèn BST thông thường.
Hàm AssignRank(root): Duyệt in-order, gán rank = cnt++.
Hàm Select(root, k): Tìm nút có rank = k.
  Nếu root->rank == k: trả về data.
  Nếu root->rank > k: tìm bên trái.
  Ngược lại: tìm bên phải.

Ví dụ mã nguồn

#include <iostream>
using namespace std;

struct Node {
    int data, rank;
    Node *left, *right;
};

Node* createNode(int d) {
    Node* n = new Node{d, 0, nullptr, nullptr};
    return n;
}

Node* insert(Node* root, int d) {
    Node* tmp = createNode(d);
    if (!root) return tmp;
    Node* t = root;
    while (t) {
        if (d < t->data) {
            if (!t->left) { t->left = tmp; break; }
            t = t->left;
        } else {
            if (!t->right) { t->right = tmp; break; }
            t = t->right;
        }
    }
    return root;
}

int cnt = 0;
void assignRank(Node* root) {
    if (root->left) assignRank(root->left);
    root->rank = cnt++;
    if (root->right) assignRank(root->right);
}

int select(Node* root, int k) {
    if (root->rank == k) return root->data;
    if (root->rank > k) return select(root->left, k);
    return select(root->right, k);
}

void display(Node* root) {
    if (root->left) display(root->left);
    cout << "data: " << root->data << " rank: " << root->rank << "\n";
    if (root->right) display(root->right);
}

int main() {
    int a[10] = {4,7,6,1,10,3,2,15,16,20};
    Node* root = nullptr;
    for (int x : a) root = insert(root, x);
    int k; cout << "Nhap k (lon thu k): "; cin >> k;
    assignRank(root);
    cout << "Rank cac nut:\n"; display(root);
    cout << "Phan tu lon thu " << k << ": " << select(root, 10-k) << endl;
    return 0;
}

Đầu ra mẫu

Nhap k: 7
Rank cac nut:
data: 1 rank: 0
data: 2 rank: 1
data: 3 rank: 2
data: 4 rank: 3
data: 6 rank: 4
data: 7 rank: 5
data: 10 rank: 6
data: 15 rank: 7
data: 16 rank: 8
data: 20 rank: 9
Phan tu lon thu 7: 4