【洛谷P3466】KLO-Building blocks

题目

题目链接：https://www.luogu.com.cn/problem/P3466
给出一个数列，将一个数加一或减一都付出\(1\)的代价。求最小代价使得数列中存在一个长度为\(m\)的子序列每一个数字都相同。

思路

很明显就是枚举每一个长度为\(m\)的区间，然后求中位数搞一搞。
求中位数可以用对顶堆。我们求出中位数后，计算小于中位数的数字之和以及大于中位数的数字之和，就可以计算出答案。
但是这样我们从区间\([i,i+m-1]\)转移到\([i+1,i+m]\)时，需要支持在堆中 删除任意元素，而这是\(STL\)自带的堆没有的功能。而又不想手写堆（
所以我们需要一个能在\(O(\log n)\)的时间复杂度内支持求最值，插入，删除的数据结构。
那么平衡树就可以了。
用平衡树代替堆，维护两棵平衡树即可。
当然这是比较暴力的方法，维护一棵平衡树也是可以的，因为平衡树支持查找中位数（大概吧）。
时间复杂度\(O(n\log m)\)。
其实写这道题就是想默一下Treap的板子233

代码

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;

const int N=800010,Inf=1e9;
int n,m,pos,v,a[N];
ll sum,ans=1e17;

struct Treenode
{
    int lc,rc,size,dat,val,cnt;
};

struct Treap
{
    Treenode tree[N];
    int tot,root;
    ll sum;
    
    int New(int x)
    {
        tree[++tot].val=x;
        tree[tot].size=tree[tot].cnt=1;
        tree[tot].dat=rand();
        return tot;
    }
    
    void pushup(int x)
    {
        tree[x].size=tree[tree[x].lc].size+tree[tree[x].rc].size+tree[x].cnt;
    }
    
    void build()
    {
        root=New(-Inf);
        tree[1].rc=New(Inf);
        pushup(1);
    }
    
    void zig(int &x)
    {
        int y=tree[x].lc;
        tree[x].lc=tree[y].rc; tree[y].rc=x; x=y;
        pushup(tree[x].rc); pushup(x);
    }
    
    void zag(int &x)
    {
        int y=tree[x].rc;
        tree[x].rc=tree[y].lc; tree[y].lc=x; x=y;
        pushup(tree[x].lc); pushup(x);
    }
    
    int Min(int x)  //求平衡树中最小值，最大值下同
    {
        if (!x) return Inf;
        if (tree[x].val==-Inf) return Min(tree[x].rc);
        return min(Min(tree[x].lc),tree[x].val);
    }
    
    int Max(int x)
    {
        if (!x) return -Inf;
        if (tree[x].val==Inf) return Max(tree[x].lc);
        return max(Max(tree[x].rc),tree[x].val);
    }
    
    void insert(int &x,int val)
    {
        if (x==root) sum+=val;
        if (!x) x=New(val);
        else if (tree[x].val==val) tree[x].cnt++;
        else if (val<tree[x].val)
        {
            insert(tree[x].lc,val);
            if (tree[tree[x].lc].dat<tree[x].dat) zig(x);
        }
        else
        {
            insert(tree[x].rc,val);
            if (tree[tree[x].rc].dat<tree[x].dat) zag(x);
        }
        pushup(x);
    }
    
    void del(int &x,int val)
    {
        if (x==root) sum-=val;
        if (tree[x].val==val)
        {
            if (tree[x].cnt>1) tree[x].cnt--;
            else if (tree[x].lc || tree[x].rc)
            {
                if (!tree[x].lc || tree[tree[x].lc].dat<tree[tree[x].rc].dat)
                    zag(x),del(tree[x].lc,val);
                else
                    zig(x),del(tree[x].rc,val);
            }
            else x=0;
        }
        else if (tree[x].val>val) del(tree[x].lc,val);
        else del(tree[x].rc,val);
        pushup(x);
    }
    
    bool find(int x,int val)  //查找平衡树中是否有val这个值
    {
        if (!x) return 0;
        if (tree[x].val==val) return 1;
        if (tree[x].val>val) return find(tree[x].lc,val);
            else return find(tree[x].rc,val);
    }
}Treap1,Treap2;

void update()  //维护两棵。。。对顶平衡树？（雾
{
    while (Treap1.Max(Treap1.root)>Treap2.Min(Treap2.root))
    {
        int x=Treap1.Max(Treap1.root);
        Treap1.del(Treap1.root,x); Treap2.insert(Treap2.root,x);
    }
    while (Treap1.tree[Treap1.root].size-Treap2.tree[Treap2.root].size>=2)
    {
        int x=Treap1.Max(Treap1.root);
        Treap1.del(Treap1.root,x); Treap2.insert(Treap2.root,x);
    }
    while (Treap2.tree[Treap2.root].size>Treap1.tree[Treap1.root].size)
    {
        int x=Treap2.Min(Treap2.root);
        Treap2.del(Treap2.root,x); Treap1.insert(Treap1.root,x);
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    srand(n*2333-m*666);  //优秀的重置种子方法
    for (int i=1;i<=n;i++)
        scanf("%d",&a[i]);
    Treap1.build(); Treap2.build();
    for (int i=1;i<m;i++)
    {
        Treap1.insert(Treap1.root,a[i]);
        update();
    }
    for (int i=m;i<=n;i++)
    {
        Treap1.insert(Treap1.root,a[i]);
        update();
        ll mid=1LL*Treap1.Max(Treap1.root);
        int cnt1=Treap1.tree[Treap1.root].size,cnt2=Treap2.tree[Treap2.root].size;
        ll cost=mid*cnt1-Treap1.sum+Treap2.sum-mid*cnt2;  //计算代价
        if (cost<ans)
            ans=cost,pos=i-m+1,v=mid;
        if (Treap1.find(Treap1.root,a[i-m+1])) Treap1.del(Treap1.root,a[i-m+1]);
            else Treap2.del(Treap2.root,a[i-m+1]);
        update();
    }
    printf("%lld\n",ans);
    for (int i=1;i<=n;i++)
        if (i>=pos && i<pos+m) printf("%d\n",v);
            else printf("%d\n",a[i]);
    return 0;
}