「模拟赛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_{v\in A,u\in A,v\not=u}f(v,u)+g(v,u)+\sum_{v\in B,u\in B,v\not = u} h(v,u)+\sum_{v\in 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\%\)的数据，\(1\leq n,m\leq 500000,0\leq |A|\leq n,1\leq a_i\leq 500000\)。

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

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

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

那这两个加起来是什么呢？就是\(C(|A|,2)+\)满足\(u\in A,v\in 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_{v\in A,u\in A,v\not=u}g(v,u)+\sum_{v\in B,u\in B,v\not = u} h(v,u)+\sum_{v\in A,u\in B}g(v,u)+h(v,u)+e(v,u) \]

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

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

然后还有一个\(e\)比较烦人，注意到若\(v\)是\(u\)的祖先，且\(v\in B,u\in 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\n", ans[x]);
	}
}