[BZOJ 1112] [POI2008] 砖块Klo 【区间K大】

题目链接：BZOJ - 1112

题目分析

枚举每一个长度为k的连续区间，求出这个区间的最优答案，更新全局答案。

可以发现，这个区间的所有柱子最终都变成这k个数的中位数时最优，那么我们就需要查询这个区间的中位数了。

找到中位数之后，我们还应该求出这个区间内小于中位数的数的和，大于中位数的数的和，从而求出操作步数。

这些需要求的值可以用线段树或平衡树来写，我写的是线段树，但是实际上这是一道POI的题目，在MAIN上的空间限制只有35MB，线段树应该是不行的。

因为平衡树只需要 O(n) 空间，所以平衡树才是正解。

代码

#include <iostream>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <cstdio>

using namespace std;

const int MaxN = 100000 + 5, MaxNode = 100000 * 20 + 15, MN = 1000000 + 5;

typedef long long LL;

int n, k, Index, Root;
int A[MaxN], T[MaxNode], Son[MaxNode][2];

const LL INF = 999999999999;

LL Ans;
LL Sum[MaxNode];

inline LL gmin(LL a, LL b) {return a < b ? a : b;}

inline void Read(int &Num)
{
	char c; c = getchar();
	while (c < '0' || c > '9') c = getchar();
	Num = c - '0'; c = getchar();
	while (c >= '0' && c <= '9')
	{
		Num = Num * 10 + c - '0';
		c = getchar();
	}
}

void Add(int &x, int s, int t, int Pos, int Num)
{
	if (x == 0) x = ++Index;
	T[x] += Num;
	Sum[x] += (LL)Pos * (LL)Num;
	if (s == t) return;
	int m = (s + t) >> 1;
	if (Pos <= m) Add(Son[x][0], s, m, Pos, Num);
	else Add(Son[x][1], m + 1, t, Pos, Num);
}

int Kth(int x, int s, int t, int k)
{
	if (s == t) return s;
	int ret, m = (s + t) >> 1;
	if (T[Son[x][0]] >= k) ret = Kth(Son[x][0], s, m, k);
	else ret = Kth(Son[x][1], m + 1, t, k - T[Son[x][0]]);
	return ret;
}

LL GetSum(int x, int s, int t, int l, int r)
{
	if (l <= s && r >= t) return Sum[x];
	int m = (s + t) >> 1;
	LL ret = 0ll;
	if (l <= m && Son[x][0]) ret += GetSum(Son[x][0], s, m, l, r);
	if (r >= m + 1 && Son[x][1]) ret += GetSum(Son[x][1], m + 1, t, l, r);
	return ret;
}

int GetNum(int x, int s, int t, int l, int r)
{
	if (l <= s && r >= t) return T[x];
	int m = (s + t) >> 1;
	int ret = 0;
	if (l <= m && Son[x][0]) ret += GetNum(Son[x][0], s, m, l, r);
	if (r >= m + 1 && Son[x][1]) ret += GetNum(Son[x][1], m + 1, t, l, r);
	return ret;
}

int main()
{
	scanf("%d%d", &n, &k);
	for (int i = 1; i <= n; ++i) Read(A[i]);
	Root = Index = 0;
	A[0] = 0;
	for (int i = 0; i <= k - 1; ++i) Add(Root, 0, MN, A[i], 1);
	Ans = INF;
	int t = k / 2 + 1, Temp;
	LL Now;
	for (int i = k; i <= n; ++i)
	{
		Add(Root, 0, MN, A[i - k], -1);
		Add(Root, 0, MN, A[i], 1);
		Temp = Kth(Root, 0, MN, t);
		Now = (LL)GetNum(Root, 0, MN, 0, Temp - 1) * (LL)Temp - GetSum(Root, 0, MN, 0, Temp - 1);
		Now += GetSum(Root, 0, MN, Temp + 1, MN) - (LL)GetNum(Root, 0, MN, Temp + 1, MN) * (LL)Temp;
		Ans = gmin(Ans, Now);
	}
	printf("%lld\n", Ans);
	return 0;
}