ZCPC17th E Easy DP Problem

ZCPC17th E Easy DP Problem

希望更丰富的阅读体验？来我的网站阅读趴！

Problem

由于这题前面的思维推到部分我没有参与，主要是现学（复习）了一下主席树，所以主要讲主席树的部分。

题目可转化为：

给一个长度为 \(n\) 的数组 \(a_i\)，有 \(q\) 个询问，每次询问区间 \([l,r]\) 中最大的 \(k\) 个数之和，再加上 \(\sum_{i=1}^{r-l+1}i^2\)。

\(n\le10^5,q\le10^5,a_i\le10^6\)

\(T\le100,\sum n\le5\times10^5,\sum q\le5\times 10^5\)

Solution

主要讲一下这题主席树的写法吧。

以前，我只知道主席树可以用来可持久化，即对于每次修改，都新建一个版本；同时，可以访问之前任意时刻的版本进行查询。

对于一个数组，使用权值线段树找最大的k个数也是很经典的线段树上二分来解决。

求最大的k个数之和，只要多维护一个 \(sum\)，在线段树上二分的时候累计答案即可。

而这题需要查询的是区间 \([l,r]\) 上的最大值。

将数组 \(a_i\) 对应到 \(root_i\)：

• \(root_0\) 代表一个全空的权值线段树
• \(root_1\) 代表只放入了 \(a_1\) 的权值线段树
• \(root_2\) 放入了 \(a_1,a_2\)
• ……

那么，想要查询区间 \([1,i]\) 上的最大k个数之和，只需要在线段树 \(root_i\) 上进行上述“线段树上二分”即可。

但是如果要查询区间 \([l,r]\) 上的最大k个数之和呢？类似前缀和与差分，只需要同时查询 \(root_{l-1}\) 和 \(root_r\) 这两颗线段树，节点的值做差即可得到区间 \([l,r]\) 上的节点权值。

Code

#define int ll
#define N 100010
#define M 1000010
int n,q;
struct node
{
	int lson, rson, cnt;
	ll sum;
#define lson(x) tree[x].lson
#define rson(x) tree[x].rson
#define cnt(x) tree[x].cnt
#define sum(x) tree[x].sum
}tree[N*20];
map<ll, ll>lsh;
ll rlsh[N];
ll lsh_cnt;
int cnt;
ll a[N];
int new_node(int old)
{
	cnt++;
	tree[cnt] = tree[old];
	return cnt;
}
void upd(int p)
{
	cnt(p) = cnt(lson(p)) + cnt(rson(p));
	sum(p) = sum(lson(p)) + sum(rson(p));
}
void build(int& p, int l, int r)
{
	p = new_node(0);
	if (l == r)
	{
		return;
	}
	int mid = (l + r) >> 1;
	build(lson(p), l, mid);
	build(rson(p), mid + 1, r);
}
void change(int& p, int p_old, int l, int r,ll v)
{
	p = new_node(p_old);
	if (l == r)
	{
		cnt(p)++;
		sum(p) += rlsh[v];
		return;
	}
	int mid = (l + r) >> 1;
	if (v <= mid)
	{
		change(lson(p), lson(p_old), l, mid, v);
	}
	else
	{
		change(rson(p), rson(p_old), mid + 1, r, v);
	}
	upd(p);
}
ll ask(int p, int pre, int l, int r,ll k)
{
	if (l == r)
	{
		return k * rlsh[l];
	}
	ll delta_cnt = cnt(rson(p)) - cnt(rson(pre));
	int mid = (l + r) >> 1;
	if (k <= delta_cnt)
	{
		return ask(rson(p), rson(pre), mid + 1, r, k);
	}
	else
	{
		return ask(lson(p), lson(pre), l, mid, k - delta_cnt) + sum(rson(p)) - sum(rson(pre));
	}
}
int root[N];
ll i2[N];
void solve() {
	cin >> n;

	for (ll i = 1; i <= n; i++) i2[i] = i2[i - 1] + i * i;
	lsh.clear();
	cnt = 0;
	lsh_cnt = 0;

	for (int i = 1; i <= n; i++)
	{
		cin >> a[i];
		lsh[a[i]] = 0;
	}
	for (auto it = lsh.begin(); it != lsh.end(); it++)
	{
		it->second = ++lsh_cnt;
		rlsh[lsh_cnt] = it->first;
	}

	build(root[0], 1, lsh_cnt);

	for (int i = 1; i <= n; i++)
	{
		change(root[i], root[i - 1], 1, lsh_cnt, lsh[a[i]]);
	}

	cin >> q;
	while (q--)
	{
		int l, r, k;
		cin >> l >> r >> k;
		ll ans = ask(root[r], root[l - 1], 1, lsh_cnt, k);
		ans += i2[r - l + 1];
		cout << ans << endl;
	}

}