「模拟赛20191019」C 推式子+贪心+树状数组

题目描述

给定一棵(n)个点的有根树，根节点编号为(1)，点有点权。

定义(d(v))表示(v)到(1)的路径上的边数。

定义(f(v,u))在(v<u)且(v)和(u)任意一个都不是另一个的祖先时为(1)，否则为(0)。

定义(g(v,u))在(v)是(u)的祖先且(v)的权值大于(u)的权值时为(1)，否则为(0)。

定义(h(v,u))在(v)是(u)的祖先且(v)的权值小于(u)的权值时为(1)，否则为(0)。

你需要将点集分成两个集合(A)和(B)，有(m)组询问，每组询问给定了集合(A)的大小，求下列表达式的最小值：

[F(A,B)=sum_{vin A,uin A,v ot=u}f(v,u)+g(v,u)+sum_{vin B,uin B,v ot = u} h(v,u)+sum_{vin A}d(v) ]

输入

第一行两个整数(n,m)。

接下来一行(n)个整数(a_i)表示第(i)个点的权值。

接下来(n−1)行，第(i)行两个整数(v,u)表示一条连接((v,u))的边。

接下来(m)行，每行一个整数表示(|A|)。

输出

共(m)行，每行一个整数表示表达式的最小值。

样例

样例输入

4 3
4 1 2 3
1 2
2 3
2 4
0
2
4

样例输出

2
2
9

数据范围

对于(100\%)的数据，(1leq n,mleq 500000,0leq |A|leq n,1leq a_ileq 500000)。

比第二题又难了不少……（差评，题目难度指数式上升）

这一坨定义看起来好难受啊，考虑转化一下。

首先看(f)的意义，发现其实就是(|A|)中互不为祖先的点对数，等价于(C(|A|,2)-)每个点在(A)中的祖先个数的和；(d)的意义是深度，也就是(A)中所有点在树上的祖先个数的和。

那这两个加起来是什么呢？就是(C(|A|,2)+)满足(uin A,vin B)且(u)是(v)的祖先的点对数量。

另外，我们定义一个新函数(e(u,v))当(u)是(v)的祖先，且(a_u=a_v)时为(1)，否则为(0)。显然([u)是(v)的祖先(])等价于(g(u,v)+h(u,v)+e(u,v))。那么化一下原式：

[F(A,B)=C(|A|,2)+sum_{vin A,uin A,v ot=u}g(v,u)+sum_{vin B,uin B,v ot = u} h(v,u)+sum_{vin A,uin B}g(v,u)+h(v,u)+e(v,u) ]

把(g)和(h)合并到前面去，于是(g)和(h)的一半部分变成全集了：

[F(A,B)=C(|A|,2)+sum_{vin A}g(v,u)+sum_{uin B} h(v,u)+sum_{vin A,uin B}e(v,u) ]

然后还有一个(e)比较烦人，注意到若(v)是(u)的祖先，且(vin B,uin A)，如果交换(v)和(u)，(g)部分会变更优，(h)部分会变更优，(e)部分也会变更优，当然会优先选择祖先……所以当一个点被加入(A)中时，他的祖先里和它权值一样的一定早就全都被加入到里面去了。所以(e)可以很轻松地计算，此时每个点加入A中带来的权值已经和其他元素毫不相干了，只需要用树状数值帮助预处理一下即可。

(Code:)

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define N 500005
#define M 1000005
#define ll long long
ll ans[N];
int n, m, A[N], tre[N], sum[N];
int tar[M], nex[M], fir[N], cnt, tim;
int fat[N], dep[N], dfn[N], out[N], val[N], idx[N];
void Read(int &p)
{
	p = 0;
	char c = getchar();
	for (; c < '0' || c > '9'; c = getchar());
	for (; c >= '0' && c <= '9'; c = getchar())p = p * 10 + c - '0';
}
void Update(int x, int v)
{
	for (int i = x; i <= N - 5; i += (i & -i))
		tre[i] += v;
}
int Getsum(int x)
{
	int ans = 0;
	for (int i = x; i; i -= (i & -i))
		ans += tre[i];
	return ans;
}
void Add(int u, int v)
{
	++cnt;
	tar[cnt] = v;
	nex[cnt] = fir[u];
	fir[u] = cnt;
}
void Dfs(int r)
{
	dfn[r] = ++tim;
	val[r] = dep[r] - Getsum(A[r]);
	Update(A[r], 1);
	for (int i = fir[r]; i; i = nex[i])
	{
		int v = tar[i];
		if (v != fat[r])
		{
			fat[v] = r;
			dep[v] = dep[r] + 1;
			Dfs(v);
		}
	}
	Update(A[r], -1);
	out[r] = tim;
};
bool cmp(int a, int b){return A[a] > A[b];}
int main()
{
	scanf("%d%d", &n, &m);
	for (int i = 1; i <= n; i++)
		Read(A[i]), idx[i] = i;
	for (int i = 1; i < n; i++)
	{
		int u, v;
		Read(u), Read(v);
		Add(u, v), Add(v, u);
	}
	Dfs(1);
	sort(idx + 1, idx + n + 1, cmp);
	for (int i = 1; i <= n; )
	{
		int x = i;
		while (x <= n && A[idx[i]] == A[idx[x]])
			x++;
		for (int k = i; k < x; k++)
			sum[idx[k]] = Getsum(out[idx[k]]) - Getsum(dfn[idx[k]] - 1);
		for (int k = i; k < x; k++)
			Update(dfn[idx[k]], 1);
		i = x;
	}
	for (int i = 1; i <= n; i++)
		val[i] -= sum[i], ans[0] += sum[i];
	sort(val + 1, val + n + 1);
	for (int i = 1; i <= n; i++)
		ans[i] = ans[i - 1] + val[i] + i - 1;
	for (; m--; )
	{
		int x;
		Read(x);
		printf("%lld
", ans[x]);
	}
}