C++和python3中的平衡搜索树介绍

1 专题介绍

本专题用来介绍平衡搜索树，并给出C++手写实现、C++容器set、python3中SortedList例子。

平衡搜索树是一种数据结构，它在树的基础上加入了平衡的特性，使得树的高度保持在一个较小的范围内，从而提高了插入、删除和查找等操作的效率。

常见的平衡搜索树包括红黑树、AVL树、Splay树等。这些树的平衡特性保证了在插入和删除元素时，树的高度保持在一个相对较小的范围内，从而保证了搜索的效率。

平衡搜索树通常具有以下特性：

1. 平衡性：树的高度保持在一个较小的范围内，使得查询、插入和删除等操作的时间复杂度在O(logN)的水平。
2. 有序性：树中元素按照一定的顺序（如升序或降序）排列，使得搜索操作更加高效。
3. 动态性：树的结构可以随时调整，以适应插入、删除等操作的变化。

平衡搜索树在许多领域都有广泛的应用，包括数据库系统、编译器设计、图形学等。它们在实现有序集合、字典等数据结构时非常有用，并且在很多算法和数据处理任务中都有重要作用。

2 初步实践

2.1 C++手写平衡搜索树：

待更新。。。

2.2 C++使用multiset：

注意，multiset中的erase(3)是删除sa中的所有元素3。

multiset中的extract(3)是删除sa中的一个元素3。

#include <iostream>
#include <set>

using namespace std;

int main() {
    multiset<int> sa; //允许重复元素的有序集合，底层通过平衡搜索树来实现
    
    //(1)插入操作
    sa.insert(3);
    sa.insert(1);
    sa.insert(4);
    
    //(2)删除操作
    sa.erase(3);
    
    //(3)查询操作
    //(3.1)查询首次大于val的元素下标
    int val = 4;
    multiset<int>::iterator it = sa.upper_bound(val);
    //(3.2)查询首次小于val的元素下标
    it = sa.lower_bound(val);
    --it;
    //(3.3)查询等于val的元素下标
    it = sa.lower_bound(val);
    if (*it != val) {
        cout << "multiset中无" << val << "!" << endl;
    }
    
    //(4)输出multiset
    for (auto it = sa.begin(); it != sa.end(); it++) {
        cout << *it << " ";
    }
    cout << endl;
    
    return 0;
}

2.3 python3使用SortedList：

from sortedcontainers import SortedList
import bisect 

sa = SortedList() #有序列表，底层用平衡搜索树实现的

#(1)插入操作
sa.add(3)
sa.add(1)
sa.add(4)

#(2)删除操作
sa.remove(3)

#(3)查询操作
##(3.1)查询变量sa中首个大于val的数的下标
val = 4
i = bisect.bisect_right(sa, val)
##(3.2)查询变量sa中首个小于val的数的下标
i = bisect.bisect_left(sa, val) - 1
##(3.3)查询变量sa中首个等于val的数的下标
i = bisect.bisect_left(sa, val)
if sa[i] != val:
    i = -1

#(4)输出有序列表
##(4.1)通过下标输出
for i in range(len(sa)):
    print(sa[i], end=" ")
print()
##(4.2)通过基于范围的for循环实现
for x in sa:
    print(x, end=" ")
print()

3 训练

题目1：100246将元素分配到两个数组中II

考点：平衡搜索树、树状数组

方法1：C++手写实现平衡搜索树

待更新。。。

在这里插入代码片

方法2：C++使用mulitset

TLE（Time Limit Exceeded）了，通过case比例为780/784。

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        int n = nums.size();
        vector<int> a = {nums[0]};
        vector<int> b = {nums[1]};
        multiset<int> sa;
        sa.insert(nums[0]);
        multiset<int> sb;
        sb.insert(nums[1]);
        
        function<bool(int)> check = [&](int val) -> bool {
            //a中大于val的元素数目
            int cnt1 = 0;
            multiset<int>::iterator it = sa.upper_bound(val);
            cnt1 = distance(it, sa.end());
            //b中大于val的元素数目
            int cnt2 = 0;
            it = sb.upper_bound(val);
            cnt2 = distance(it, sb.end());
            
            if (cnt1 > cnt2) return true;
            else if (cnt1 < cnt2) return false;
            else {
                if (a.size() <= b.size()) return true;
                else return false;
            }
                
            
            cout << "-1" << endl;
            return true;
        };
        
        for (int i = 2; i < n; ++i) {
            //val时刻在变化
            if (check(nums[i])) {
                a.emplace_back(nums[i]);
                sa.insert(nums[i]);
            } else {
                b.emplace_back(nums[i]);
                sb.insert(nums[i]);
            }
        }
        
        vector<int> res = a;
        res.insert(res.end(), b.begin(), b.end());
        return res;
    }
};

方法3：python3使用sortedcontainers中的SortedList

from sortedcontainers import SortedList

class Solution:
    def resultArray(self, nums: List[int]) -> List[int]:
        n = len(nums)
        a = [nums[0]]
        b = [nums[1]]
        
        sa = SortedList()
        sa.add(nums[0])
        sb = SortedList()
        sb.add(nums[1])
        
        for i in range(2, n):
            
            def check(val: int) -> bool:
                nonlocal a, b
                nonlocal sa, sb
                
                cnt1 = len(sa) - bisect.bisect_right(sa, val)
                cnt2 = len(sb) - bisect.bisect_right(sb, val)
                
                if cnt1 > cnt2:
                    return True
                elif cnt1 < cnt2:
                    return False
                else:
                    if len(a) <= len(b):
                        return True
                    else:
                        return False
                print("-1")
                return True
            
            
            if check(nums[i]):
                a.append(nums[i])
                sa.add(nums[i])
            else:
                b.append(nums[i])
                sb.add(nums[i])
        
        return a + b

方法4：使用树状数组，

class Fenwick {
    vector<int> tree;
public:
    Fenwick(int n) : tree(n) {}

    //把下标为i的元素增加1
    void add(int i) {
        while (i < tree.size()) {
            tree[i]++;
            i += i & -i;
        }
    }

    //返回下标在[1,i]的元素之和
    int pre(int i) {
        int res = 0;
        while (i > 0) {
            res += tree[i];
            i &= i - 1;
        }
        return res;
    }
};

class Solution {
public:
    vector<int> resultArray(vector<int>& nums) {
        auto sorted = nums;
        ranges::sort(sorted);
        sorted.erase(unique(sorted.begin(), sorted.end()), sorted.end());
        int m = sorted.size();

        vector<int> a{nums[0]}, b{nums[1]};
        Fenwick t1(m + 1), t2(m + 1);
        t1.add(ranges::lower_bound(sorted, nums[0]) - sorted.begin() + 1);
        t2.add(ranges::lower_bound(sorted, nums[1]) - sorted.begin() + 1);
        for (int i = 2; i < nums.size(); ++i) {
            int x = nums[i];
            int v = ranges::lower_bound(sorted, x) - sorted.begin() + 1;
            int gc1 = a.size() - t1.pre(v);
            int gc2 = b.size() - t2.pre(v);
            if (gc1 > gc2 || gc1 == gc2 && a.size() <= b.size()) {
                a.push_back(x);
                t1.add(v);
            } else {
                b.push_back(x);
                t2.add(v);
            }
        }
        a.insert(a.end(), b.begin(), b.end());
        return a;
    }
};