JZOJ 1078. 【GDOI2006】The Kth Element

\(\text{Problem}\)

给定一个整数序列 \(a[1..N]\)，定义 \(sum[i][j]=a[i]+a[i+1]+...+a[j]\)，将所有的 \(sum[i][j]\) 从小到大排序(其中 \(i,j\) 满足 \(1<=i<=j<=N\) )，得到一个长为 \(N*(N+1)/2\) 的序列，求该序列中的第 \(K\) 个元素。

\(\text{Solution}\)

非常经典的题

二分答案然后 \(check\)

\(\text{Code}\)

#include <cstdio>
#include <algorithm>
#include <iostream>
#define ls (p << 1)
#define rs (ls | 1)
#define RE register
#define IN inline
using namespace std;
typedef long long LL;

const int N = 2e4 + 5;
const LL INF = 1e18;
int n, m, len;
LL a[N], b[N], sum[N * 4];

void build(int p, int l, int r)
{
	sum[p] = 0;
	if (l == r) return;
	int mid = l + r >> 1;
	build(ls, l, mid), build(rs, mid + 1, r);
}
void Modify(int p, int l, int r, int x, int v)
{
	if (x > r || x < l) return;
	if (l == r) return sum[p] += v, void();
	int mid = l + r >> 1;
	if (x <= mid) Modify(ls, l, mid, x, v);
	else Modify(rs, mid + 1, r, x, v);
	sum[p] = sum[ls] + sum[rs];
}
LL Query(int p, int l, int r, int x)
{
	if (x > r || x < l) return 0;
	if (l == r) return sum[p];
	int mid = l + r >> 1;
	if (x <= mid) return Query(ls, l, mid, x) + sum[rs];
	return Query(rs, mid + 1, r, x);
}

IN int check(LL mid)
{
	int res = 0;
	build(1, 1, len), Modify(1, 1, len, lower_bound(b + 1, b + len + 1, 0) - b, 1);
	for(RE int i = 1; i <= n; i++)
	{
		int x = lower_bound(b + 1, b + len + 1, a[i] - mid) - b;
		res += Query(1, 1, len, x);
		Modify(1, 1, len, lower_bound(b + 1, b + len + 1, a[i]) - b, 1);
	}
	return res >= m;
}

int main()
{
	scanf("%d%d", &n, &m);
	for(RE int i = 1; i <= n; i++) scanf("%lld", &a[i]), a[i] += a[i - 1], b[i] = a[i];
	b[n + 1] = 0, sort(b + 1, b + n + 2), len = unique(b + 1, b + n + 2) - b - 1;
	LL l = -INF, r = INF, mid, ans;
	while (l <= r)
	{
		mid = (l + r) >> 1;
		if (check(mid)) ans = mid, r = mid - 1;
		else l = mid + 1;
	}
	printf("%lld\n", ans);
}