【巧用set实现对有序数组O(logn)时间复杂度增、删、查、改、二分操作】codeforces 1041 C. Coffee Break

题意

第一行输入三个整数 \(n,m,d(1 \leq n \leq 2 * 10^5, n \leq m \leq 10^9, 1 \leq d \leq n)\)，第二行输入 \(n\) 个整数，保证每个数均不大于 \(m\)。

在每一天你都可以任意选择一个未选过的数 \(a_i\)，随后可以继续选任意一个大于 \(a_i + d\) 的数 \(a_j\)；接下来可以再选任意一个大于 \(a_j + d\) 的数 \(a_k\)；最后重复执行上述步骤。

问：最少需要多少天，可以选完全部的数？

题解

先给结论：先从小到大排序，然后选中剩下的数组的第一个数 \(a_0\)，随后选中数组中第一个满足 \(a_i \geq a_0 + d + 1\) 的数 \(a_i\)，下一步再选中数组中第一个满足 \(a_j \geq a_i + d\) 的数 \(a_j\)，形成的子数组数目便是答案。

疑问解析：为什么优先选择最小的满足 \(a_j \geq a_i + d + 1\)的数是正确的？

答：不妨假设数组排序后为 \(a_0, a_1, ..., a_i, a_j, ..., a_n\)，且 \(a_i\) 是最小的满足大于 \(a_0 + d\) 的数。那么明显在选中 \(a_0\) 的下一步，既可以选择 \(a_i\)，也可以选择 \(a_j\)。此时若 \(a_j - a_i \leq d\)，那明显二者至多只能够选其一，那么无论选择哪一个，都必然意味着需要另外再选择出一个子数组。但如果 \(a_j - a_i \geq d + 1\)，就可以先选 \(a_i\) 而后一并选上 \(a_j\)。因此，先选较小者更优。

接下来分析具体代码实现：

如何在数组中找出第一个比选中的数大 \(d\) 以上的数呢？

明显可以线性查找，但是当数组中任意两个数的差值均小于 \(d\) 时，时间复杂度为 \(O(n^2)\)，明显不符合要求

观察我们的前置条件排序，所以可以考虑使用二分法来使得查找的时间复杂度降低为 \(O(logn)\)

如何选择过的数从数组中移除呢？

若通过类似插入排序等思路，时间复杂度为 \(O(n^2)\)，明显不符合要求

此时思考：既要排序，又要可以二分，还要支持快速删除，我们可以联想到红黑树的性质，但是手撕红黑树太硬核了，可以借助 \(set\) 或 \(map\) 实现，这二者的查找/插入/删除/修改操作时间复杂度都是 \(O(logn)\)，符合时间复杂度要求。

参考代码

#include<bits/stdc++.h>
#define IOS ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr);
using namespace std;

constexpr int INF = 0x3f3f3f3f;
constexpr int N = 2e5 + 7;
int n, m, d, cnt, a;
int ans[N];

int main() {
    IOS
    set<PII> st;
    cin >> n >> m >> d;
    for (int i = 0; i < n; ++ i) {
        cin >> a;
        st.insert(PII(a, i));
    }
    while (!st.empty()) {
        auto f = st.begin();
        int key = (*f).first;
        ans[(*f).second] = ++ cnt;
        st.erase(f);
        while (!st.empty()) {
            auto ne = st.upper_bound(PII(key + d, INF));
            if (ne != st.end()) {
                key = (*ne).first;
                ans[(*ne).second] = cnt;
                st.erase(ne);
            } else break;
        }
    }
    cout << cnt << '\n';
    for (int i = 0; i < n; ++ i) cout << ans[i] << ' ';
    return 0;
}