题解 | #【模板】集合操作#

题目链接

【模板】集合操作

题目描述

你需要动态维护一个初始为空的集合 M，支持以下几种操作：

插入: 向集合 M 中插入一个数 x。
删除: 从集合 M 中删除一个数 x。
查询存在性: 查询数 x 是否在集合 M 中。
查询大小: 查询集合 M 中元素的个数。
查询前驱: 查询 x 的前驱（小于 x 的最大数）。
查询后继: 查询 x 的后继（大于 x 的最小数）。

输入描述:

第一行一个整数 q，表示操作总数。
接下来 q 行，每行输入操作类型 op 和相应的操作数。

输出描述:

对于查询操作类型 3，如果数存在，输出 "YES"，否则输出 "NO"。
对于查询操作类型 4, 5, 6，输出相应的结果，每个结果占一行。
如果前驱或后继不存在，输出 -1。

示例输入:

输出:

YES
10
20

解题思路

这道题是典型的核心代码模式，要求我们填充预定义的函数来实现对一个集合的多种操作。核心在于选择一个高效的数据结构来支持插入、删除、查找、查前驱/后继等功能。

核心数据结构：平衡二叉搜索树（Balanced Binary Search Tree） 是最理想的选择。它能保持元素有序，并确保所有操作的时间复杂度都在对数级别。
- C++：std::set 是基于红黑树（一种平衡二叉搜索树）实现的，提供了我们需要的所有接口。
- Java：java.util.TreeSet 同样是基于红黑树实现的，功能完备。
- Python：标准库没有内置的平衡二叉搜索树。可以使用 list 和 bisect 模块模拟，但插入/删除是 $\mathcal{O}(N)$ 复杂度，可能会超时。我们将把 list 定义在全局，供各个函数操作。
函数实现:
- 全局变量: 在 C++ 和 Java 中，我们将 set/TreeSet 定义为静态全局变量。在 Python 中，将 list 定义为全局变量。
- 插入/删除/查找/大小: 直接调用相应数据结构的 insert/add、erase/remove、find/contains (count in C++)、size 方法。
- 前驱查询 (getPre): 找到小于 x 的最大元素。在 C++ 中用 lower_bound 然后回退一步，在 Java 中用 lower()。如果找不到，返回 -1。
- 后继查询 (getBack): 找到大于 x 的最小元素。在 C++ 中用 upper_bound，在 Java 中用 higher()。如果找不到，返回 -1。

我们只需要将这些逻辑填充到题目给定的函数模板中即可。

代码

cpp
java
python

#include<bits/stdc++.h>
using namespace std;

set<int> s;

void insertValue(int x){
    s.insert(x);
}
void eraseValue(int x){
    s.erase(x);
}
int xInSet(int x){
    return s.count(x);
}
int sizeOfSet(){
    return s.size();
}
int getPre(int x){
    auto it = s.lower_bound(x);
    if(it == s.begin()) return -1;
    it--;
    return *it;
}
int getBack(int x){
    auto it = s.upper_bound(x);
    if(it == s.end()) return -1;
    return *it;
}

int main(){
    int q,op,x;
    cin>>q;
    while(q--){
        cin>>op;
        if(op==1){
            cin>>x;
            insertValue(x);
        }
        if(op==2){
            cin>>x;
            eraseValue(x);
        }
        if(op==3){
            cin>>x;
            if(xInSet(x)){
                cout<<"YES\n";
            }else{
                cout<<"NO\n";
            }
        }
        if(op==4){
            cout<<sizeOfSet()<<endl;
        }
        if(op==5){
            cin>>x;
            cout<<getPre(x)<<endl;
        }
        if(op==6){
            cin>>x;
            cout<<getBack(x)<<endl;
        }
    }
    return 0;
}

import java.util.*;

public class Main {
    static TreeSet<Integer> set = new TreeSet<>();

    public static void insertValue(int x) {
        set.add(x);
    }

    public static void eraseValue(int x) {
        set.remove(x);
    }

    public static boolean xInSet(int x) {
        return set.contains(x);
    }

    public static int sizeOfSet() {
        return set.size();
    }

    public static int getPre(int x) {
        Integer pre = set.lower(x);
        return (pre == null) ? -1 : pre;
    }

    public static int getBack(int x) {
        Integer post = set.higher(x);
        return (post == null) ? -1 : post;
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        int q = scanner.nextInt();
        while (q-- > 0) {
            int op = scanner.nextInt();
            switch (op) {
                case 1:
                    insertValue(scanner.nextInt());
                    break;
                case 2:
                    eraseValue(scanner.nextInt());
                    break;
                case 3:
                    System.out.println(xInSet(scanner.nextInt()) ? "YES" : "NO");
                    break;
                case 4:
                    System.out.println(sizeOfSet());
                    break;
                case 5:
                    System.out.println(getPre(scanner.nextInt()));
                    break;
                case 6:
                    System.out.println(getBack(scanner.nextInt()));
                    break;
            }
        }
        scanner.close();
    }
}

import sys
import bisect

# 注意: Python 标准库没有平衡二叉搜索树，使用 list+bisect 模拟。
# 对于大数据量的插入/删除操作，此代码可能会超时 (Time Limit Exceeded)。
data = []

def insertValue(x):
    idx = bisect.bisect_left(data, x)
    if idx == len(data) or data[idx] != x:
        data.insert(idx, x) # O(N)

def eraseValue(x):
    idx = bisect.bisect_left(data, x)
    if idx < len(data) and data[idx] == x:
        data.pop(idx) # O(N)

def xInSet(x):
    idx = bisect.bisect_left(data, x)
    return idx < len(data) and data[idx] == x

def sizeOfSet():
    return len(data)

def getPre(x):
    idx = bisect.bisect_left(data, x)
    if idx > 0:
        return data[idx - 1]
    else:
        return -1

def getBack(x):
    idx = bisect.bisect_right(data, x)
    if idx < len(data):
        return data[idx]
    else:
        return -1

def main():
    # 使用 sys.stdin.readline() 提高 I/O 效率
    # 模板中的 input() 可能会在数据量大时导致 TLE
    q = int(sys.stdin.readline())
    for _ in range(q):
        line = list(map(int, sys.stdin.readline().split()))
        op = line[0]
        
        if op == 1:
            insertValue(line[1])
        elif op == 2:
            eraseValue(line[1])
        elif op == 3:
            print("YES" if xInSet(line[1]) else "NO")
        elif op == 4:
            print(sizeOfSet())
        elif op == 5:
            print(getPre(line[1]))
        elif op == 6:
            print(getBack(line[1]))

if __name__ == "__main__":
    main()

算法及复杂度

算法：平衡二叉搜索树 (C++/Java) / 有序列表+二分查找 (Python)
时间复杂度:
- C++/Java: 所有操作（插入、删除、查找、前驱、后继）的平均和最坏时间复杂度均为 $\mathcal{O}(\log N)$ ，其中 N 是集合中的元素个数。size() 操作为 $\mathcal{O}(1)$ 。处理 Q 个操作的总时间复杂度为 $\mathcal{O}(Q \log N_{max})$ 。
- Python: 查找、前驱、后继操作为 $\mathcal{O}(\log N)$ 。但插入和删除为 $\mathcal{O}(N)$ 。总时间复杂度最坏情况下为 $\mathcal{O}(Q \cdot N_{max})$ 。
空间复杂度： $\mathcal{O}(N_{max})$ ，其中 $N_{max}$ 是集合中存储过的最大元素数量，用于存储集合本身。