题目描述
给定一棵\(n\)个点的有根树，根节点编号为\(1\)，点有点权。
定义\(d(v)\)表示\(v\)到\(1\)的路径上的边数。
定义\(f(v,u)\)在\(v
定义\(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
#include
#include
#include
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]);
}
}