「模拟赛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$行，每行一个整数表示表达式的最小值。

样例

样例输入

样例输出

数据范围

对于$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]);

	}

}