题目描述

难度分:$2100$

输入$n$、$k(1 \leq k \leq n \leq 5 \times 10^5)$、$d(0 \leq d \leq 10^9)$和长为$n$的数组$a(1 \leq a[i] \leq 10^9)$。

你需要把$a$分成若干个集合，每个集合的大小至少为$k$，且同一集合中任意两数之差至多为$d$。

能否做到？输出YES或NO。

输入样例$1$

6 3 10
7 2 7 7 4 2

输出样例$1$

YES

输入样例$2$

6 2 3
4 5 3 13 4 10

输出样例$2$

YES

输入样例$3$

3 2 5
10 16 22

输出样例$3$

NO

算法

线段树优化DP

先贪一下，集合内任意两数之差的绝对值不超过$d$，那只要极差不超过$d$就好了。对数组排序，然后从后往前遍历，让$a[i]$成为某个集合的开头。

状态定义

$f[i]$表示后缀$[i,n]$能否被划分为大小至少为$k$（取值为$1$表示可以，为$0$表示不可以），且子段内极差$\leq d$的若干段。在这个定义下，答案就应该是$f[1]$。

状态转移

初始化$f[n+1]=1$，这是$base$ $case$。

如果要以$i$为集合的第一个元素，那么可以通过二分找到下一个集合的第一个元素，也就是$a[j]-a[i] \gt d$的最小$j$，这是下一个集合首元素的最远位置，再远就会使得当前集合$[i,j)$的极差超过$d$。而每个集合的大小至少要是$k$，因此这个$j$还可以往左移动到$i+k$（如果$j-i \lt k$就无解，$f[i]=0$）。

所以只要$f[i+k:j]$中存在一个值为$1$，$f[i]=1$就成立，所以状态转移方程为

$f[i]=max_{index \in [i+k,j]}f[index]$

此时发现这是一个子数组动态求最值的操作，用线段树存储DP数组就能在$O(log_2n)$的时间复杂度下完成转移。

复杂度分析

时间复杂度

状态数量是$O(n)$，单次转移的时间复杂度为$O(log_2n)$，整个算法的时间复杂度为$O(nlog_2n)$。

空间复杂度

除开输入的原始数组$a$，空间消耗的瓶颈就是线段树的空间，额外空间复杂度为$O(n)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;

const int N = 500010;
int n, k, d, a[N];

class SegmentTree {
public:
    struct Info {
        int l, r, v, maximum;
        Info() {}
        Info(int left, int right, int val): l(left), r(right), v(val), maximum(val) {}
    } seg[N<<2];

    explicit SegmentTree() {}

    void build(int u, int l, int r) {
        if(l == r) {
            seg[u] = Info(l, r, -1);
        }else {
            int mid = l + r >> 1;
            build(u<<1, l, mid);
            build(u<<1|1, mid + 1, r);
            pushup(u);
        }
    }

    void modify(int pos, int val) {
        modify(1, pos, val);
    }

    Info query(int l, int r) {
        return query(1, l, r);
    }

private:
    void modify(int u, int pos, int val) {
        if(seg[u].l == pos && seg[u].r == pos) {
            seg[u] = Info(pos, pos, val);
        }else {
            int mid = seg[u].l + seg[u].r >> 1;
            if(pos <= mid) {
                modify(u<<1, pos, val);
            }else {
                modify(u<<1|1, pos, val);
            }
            pushup(u);
        }
    }

    Info query(int u, int l, int r) {
        if(l <= seg[u].l && seg[u].r <= r) return seg[u];
        int mid = seg[u].l + seg[u].r >> 1;
        if(r <= mid) {
            return query(u<<1, l, r);
        }else if(mid < l) {
            return query(u<<1|1, l, r);
        }else {
            return merge(query(u<<1, l, r), query(u<<1|1, l, r));
        }
    }

    void pushup(int u) {
        seg[u] = merge(seg[u<<1], seg[u<<1|1]);
    }

    Info merge(const Info& lchild, const Info& rchild) {
        Info info;
        info.l = lchild.l;
        info.r = rchild.r;
        info.v = lchild.v + rchild.v;
        info.maximum = max(lchild.maximum, rchild.maximum);
        return info;
    }
};

int main() {
    scanf("%d%d%d", &n, &k, &d);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
    }
    sort(a + 1, a + n + 1);
    SegmentTree seg;
    seg.build(1, 1, n + 1);
    seg.modify(n + 1, 1);
    for(int i = n; i >= 1; i--) {
        int j = upper_bound(a + 1, a + n + 1, a[i] + d) - a;
        int cnt = j - i;
        if(cnt < k) {
            seg.modify(i, 0);
        }else {
            int res = seg.query(i + k, j).maximum;
            if(res > 0) {
                seg.modify(i, 1);
            }else {
                seg.modify(i, 0);
            }
        }
    }
    puts(seg.query(1, 1).v == 1? "YES": "NO");
    return 0;
}