P3466 [POI 2008] KLO-Building blocks

看题解都六七年前的了，而且没有用 set 做的，于是就来写一篇，过不过随缘吧。

题意

• 给定一个 $n$ 个元素的序列 $h_1,h_2,\ldots,h_n$。
• 每次操作，可以对任意 $h_i$ 加 $1$ 或减 $1$。
• 要使长为 $k$ 的连续子序列相同，求最少操作次数。

分析

对于区间 $[i, i + k - 1]$，若将区间内所有元素都设置为 $x$，则总操作次数为：

当 $x$ 为 $h_i,h_{i+1}, \ldots, h_{i + k - 1}$ 的中位数时，总操作次数最少，即中位数贪心。关于证明，可以参考这篇文章。

在区间滑动时，需要动态维护区间内的中位数。考虑使用两个数据结构分别维护区间内小于等于中位数的元素以及大于中位数的元素，则两个数据结构内元素个数差应小于等于 $1$，所以需要实时维护。

这个数据结构需要支持插入、删除、查询最大最小值的操作。所以 std::multiset 是个很好的选择（std::set 会去重）。它能在 $O(\log n)$ 复杂度内进行插入、删除并且自动排序。

推算

初始时，$h_1$ 为中位数。对于后续元素若 $\le$ 中位数，则加到左侧，否则加到右侧。若 $i > k$，还需要删除元素 $h_{i - k}$。在插入、删除操作后，需要平衡左右两侧的元素数量。

以样例序列 $3, 9, 2, 3, 1$ 为例，将中位数存在左侧最后一个元素：

初始时中位数为 $3$。

l: 3(中位数)
r: 

插入 $9$，由于 $9 > 3$，所以 $9$ 插入右侧。

l: 3(中位数)
r: 9

插入 $2$，由于 $2 < 3$，所以 $2$ 插入左侧。

l: 2, 3(中位数)
r: 9

此时将 $3,9,2$ 设置为中位数 $3$, 需要 $\lvert 3 - 3 \rvert + \lvert 9 - 3 \rvert + \lvert 2 - 3 \rvert = 7$ 次操作。

插入 $3$，同时由于 $k = 3$，所以要删除 $h_1 = 3$，所以没有变化。

l: 2, 3(中位数)
r: 9

此时，还是需要 $7$ 次操作。

插入 $1$，由于 $1 < 3$，所以 $1$ 插入左侧。

l: 1, 2, 3(中位数)
r: 9

由于左侧元素数量为 $3$，右侧为 $1$，所以要将左侧的最大值，即 $3$ 移到右侧。新的中位数为 $2$。

l: 1, 2(中位数)
r: 3, 9

同时，要删除 $9$。

l: 1, 2(中位数)
r: 3

此时，将 $2,3,1$ 设置为中位数 $2$ 需要 $\lvert 2 - 2 \rvert + \lvert 3 - 2 \rvert + \lvert 1 - 2 \rvert = 2$ 次操作。

即最终答案为 $2$。

实现

用两个 std::multiset l, r 代表左侧与右侧，用 suml 与 sumr 记录左侧和与右侧和。无论是插入还是删除，都需要平衡两侧的元素个数。

对于每一个区间，记此时的中位数为 $c$，左侧元素个数为 $szl$, 右侧元素个数为 $szr$，那么操作次数应为 $(szl \cdot c - suml) + (sumr - szr \cdot c)$。

std::multiset 的插入删除时间复杂度均为 $O(\log n)$，同一时间内左侧右侧元素和最多有 $k$ 个，所以总时间复杂度为 $O(n\log k)$。

代码

#include <bits/stdc++.h>
using namespace std;
#define seta(f, x) memset(f, x, sizeof f)
using ll = long long;
ll read() {
    ll x = 0, f = 1; char ch = getchar();
    while (ch < '0' || ch > '9') {
        if (ch == '-') {
            f = -1;
        }
        ch = getchar();
    }
    while (ch >= '0' && ch <= '9') {
        x = x * 10 + ch - '0', ch = getchar();
    }
    return x * f;
}
#define rd read()
void write(ll x) {
    if (x < 0) {
        putchar('-'), x = -x;
    } 
    if (x > 9) {
        write(x / 10);
    }
    putchar(x % 10 + '0');
}
#define wln(x) write(x); putchar('\n');
#define wt(x) write(x); putchar(' ');
#define db(x) printf(#x); putchar(':'); wln(x);

const int N = 1e5 + 10;
int n, k;
ll ans = 9e18;
multiset<int> l, r;
// l 比中位数小的，r 比中位数大的
ll suml, sumr; 
int pos, cen;
// 最小方案的是将 [pos, pos + k) 的数都重置为 cen.
// suml 左边和，sumr 右边和，*l.rbegin() 中位数
int a[N];

int main() {
    n = rd, k = rd;
    for (int i = 1; i <= n; i++) {
    	a[i] = rd;
    }
    for (int i = 1, x = a[i]; i <= n; i++, x = a[i]) {
    	if (l.empty() || x <= *l.rbegin()) {
    		l.insert(x);
    		suml += x;
    	} else {
    		r.insert(x);
    		sumr += x;
    	}
    	while (l.size() > r.size() + 1) { // 平衡两侧
    		r.insert(*l.rbegin());
    		sumr += *l.rbegin();
    		suml -= *l.rbegin();
    		l.erase(prev(l.end()));
    	}
    	while (l.size() < r.size()) {
    		l.insert(*r.begin());
    		suml += *r.begin();
    		sumr -= *r.begin();
    		r.erase(r.begin());
    	}
    	if (i > k) {
    		auto it = l.find(a[i - k]);
    		if (it != l.end()) {
    			l.erase(it);
    			suml -= a[i - k];
    		} else {
    			it = r.find(a[i - k]);
    			r.erase(it);
    			sumr -= a[i - k];
    		}    	
    		while (l.size() > r.size() + 1) {
	    		r.insert(*l.rbegin());
	    		sumr += *l.rbegin();
	    		suml -= *l.rbegin();
	    		l.erase(prev(l.end()));
	    	}
	    	while (l.size() < r.size()) {
	    		l.insert(*r.begin());
	    		suml += *r.begin();
	    		sumr -= *r.begin();
	    		r.erase(r.begin());
	    	}
    	}
    	if (i >= k) {
			ll cur = (l.size() * (*l.rbegin()) - suml + sumr - r.size() * (*l.rbegin()));
			if (ans > cur) {
				cen = *l.rbegin();
				pos = i - k + 1;
				ans = cur;
			}  		
    	}
    }
    wln(ans);
    for (int i = 1; i <= n; i++) {
    	if (i < pos || i >= pos + k) {
    		wln(a[i]);
    	} else {
    		wln(cen);
    	}
    }
    return 0;
}

AC Record

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N=2e5+10;
#define lc(p) tr[p].l
#define rc(p) tr[p].r
mt19937 rnd(114514);
struct node{int l,r,v,s,siz,k;}tr[N];int trlen,rt;
int newd(int v){tr[++trlen]={0,0,v,v,1,rnd()};return trlen;}
void pushup(int p)
{
	tr[p].s=tr[lc(p)].s+tr[rc(p)].s+tr[p].v;
	tr[p].siz=tr[lc(p)].siz+tr[rc(p)].siz+1;
}
void split(int p,int v,int &x,int &y)
{
	if(!p){x=y=0;return;}
	if(tr[p].v<=v)
	{
		x=p;
		split(rc(p),v,rc(x),y);
	}
	else
	{
		y=p;
		split(lc(p),v,x,lc(y));
	}
	pushup(p);
}
int merge(int x,int y)
{
	if(!x||!y)return x+y;
	if(tr[x].k<tr[y].k)
	{
		rc(x)=merge(rc(x),y);
		pushup(x);
		return x;
	}
	else
	{
		lc(y)=merge(x,lc(y));
		pushup(y);
		return y;
	}
}
void add(int v)
{
	int x,y,z;
	split(rt,v-1,x,y);
	z=newd(v);
	rt=merge(merge(x,z),y);
}
void del(int v)
{
	int x,y,z;
	split(rt,v-1,x,y);
	split(y,v,y,z);
	y=merge(lc(y),rc(y));
	rt=merge(merge(x,y),z);
}
int getk(int p,int k)
{
	if(!p)return 0;
	if(k<=tr[lc(p)].siz)
		return getk(lc(p),k);
	if(k==tr[lc(p)].siz+1)
		return p;
	return getk(rc(p),k-tr[lc(p)].siz-1);
}
int a[N];
signed main()
{
	int n,k;cin>>n>>k;
	int mx,id,res;mx=1e18;
	for(int i=1;i<=n;i++)
	{
		cin>>a[i];
		add(a[i]);if(i>k)del(a[i-k]);
		if(i<k)continue;
		int mid=getk(rt,(k+1)/2),x,y,sum1,sum2,siz1,siz2;
		mid=tr[mid].v;
//		cout<<"v:"<<mid<<'\n';
		split(rt,mid-1,x,y);
		sum1=tr[x].s,siz1=tr[x].siz;
		rt=merge(x,y);
		split(rt,mid,x,y);
		sum2=tr[y].s,siz2=tr[y].siz;
		rt=merge(x,y);
//		cout<<sum1<<' '<<siz1<<' '<<sum2<<' '<<siz2<<'\n';
		int sum=sum2-siz2*mid+siz1*mid-sum1;
//		cout<<sum<<'\n';
		if(sum<mx)mx=sum,id=i,res=mid;
	}
	for(int i=id-k+1;i<=id;i++)a[i]=res;
	cout<<mx<<'\n';
	for(int i=1;i<=n;i++)cout<<a[i]<<'\n';
	return 0;
}