【洛谷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;
}