如何在C++的Order Statistic Tree中支持重复元素？

默认的__gnu_pbds有序统计树（order statistic set）基于唯一键设计，无法直接存储重复元素。针对你的需求，这里提供两种可行方案：

方案一：用Pair包装键（最简单实现） #

核心思路是给每个重复键绑定一个唯一标识，让整个pair成为唯一键，插入时维护一个递增计数器即可。查询find_by_order时只需取pair的第一个元素。

代码示例 #

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <iostream>
using namespace std;
using namespace __gnu_pbds;

// 定义有序统计树，键为pair<int, int>：第一个元素是实际值，第二个是唯一标识
typedef tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, tree_order_statistics_node_update> ost_tree;

int main() {
    ost_tree s;
    int unique_id = 0;

    // 插入目标集合{1,1,2,2,3,4}
    s.insert({1, unique_id++});
    s.insert({1, unique_id++});
    s.insert({2, unique_id++});
    s.insert({2, unique_id++});
    s.insert({3, unique_id++});
    s.insert({4, unique_id++});

    // 验证需求：第0、1位均为1
    cout << s.find_by_order(0)->first << endl; // 输出1
    cout << s.find_by_order(1)->first << endl; // 输出1

    // 额外示例：查询小于2的元素个数
    cout << s.order_of_key({2, 0}) << endl; // 输出2
    return 0;
}

优缺点 #

• 优点：完全复用现有API，无需修改树结构，代码量小，易维护。
• 缺点：每个重复元素都作为独立节点存储，内存占用略高；order_of_key查询时需传入{key, 0}来获取该键的起始排名。

方案二：自定义节点存储计数（内存高效） #

如果追求内存效率，可以让每个节点存储键的出现次数，同时维护子树的总元素数。这需要自定义tree_node_update策略来重写统计逻辑。

代码示例 #

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <iostream>
using namespace std;
using namespace __gnu_pbds;

// 自定义节点更新策略：维护子树总元素数
template<typename Node_CItr, typename Node_Itr, typename Cmp_Fn, typename _Alloc>
struct count_order_statistics_update {
    typedef int metadata_type; // 存储子树总元素数

    // 更新节点的元数据（子树大小）
    void operator()(Node_Itr node, Node_CItr end_node) {
        int total = node->m_value.second; // 当前键的计数
        if (node->m_left != end_node) total += node->m_left->get_metadata();
        if (node->m_right != end_node) total += node->m_right->get_metadata();
        node->set_metadata(total);
    }

    // 查找第k个元素
    Node_CItr find_by_order(Node_CItr node, Node_CItr end_node, size_t k) const {
        while (node != end_node) {
            size_t left_size = (node->m_left != end_node) ? node->m_left->get_metadata() : 0;
            if (k < left_size) {
                node = node->m_left;
            } else if (k < left_size + node->m_value.second) {
                return node;
            } else {
                k -= left_size + node->m_value.second;
                node = node->m_right;
            }
        }
        return end_node;
    }

    // 计算小于目标键的元素总数
    size_t order_of_key(Node_CItr node, Node_CItr end_node, const pair<int, int>& key) const {
        size_t count = 0;
        while (node != end_node) {
            if (Cmp_Fn()(node->m_value.first, key.first)) {
                count += (node->m_left != end_node) ? node->m_left->get_metadata() : 0;
                count += node->m_value.second;
                node = node->m_right;
            } else {
                node = node->m_left;
            }
        }
        return count;
    }
};

// 定义树：键为pair<int, int>，第一个是实际值，第二个是计数
typedef tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, count_order_statistics_update> count_ost_tree;

// 封装插入操作：存在则计数+1，否则插入新节点
void insert(count_ost_tree& tree, int key) {
    auto it = tree.find({key, 0});
    if (it != tree.end()) {
        int cnt = it->second;
        tree.erase(it);
        tree.insert({key, cnt + 1});
    } else {
        tree.insert({key, 1});
    }
}

// 封装删除操作：计数-1，为0则删除节点
void erase(count_ost_tree& tree, int key) {
    auto it = tree.find({key, 0});
    if (it != tree.end()) {
        int cnt = it->second;
        tree.erase(it);
        if (cnt > 1) tree.insert({key, cnt - 1});
    }
}

int main() {
    count_ost_tree s;
    insert(s, 1);
    insert(s, 1);
    insert(s, 2);
    insert(s, 2);
    insert(s, 3);
    insert(s, 4);

    // 验证需求
    cout << s.find_by_order(0)->first << endl; // 输出1
    cout << s.find_by_order(1)->first << endl; // 输出1

    // 查询小于2的元素个数
    cout << s.order_of_key({2, 0}) << endl; // 输出2
    return 0;
}

优缺点 #

• 优点：重复元素仅存储一个节点加计数，内存占用低；统计逻辑更贴合实际需求。
• 缺点：需要自定义节点更新策略，对__gnu_pbds内部结构有一定了解；插入/删除需要封装，不能直接用原生insert/erase。

内容的提问来源于stack exchange，提问作者Roy