HDU7458 旅行 题解

前言

感觉是非常优秀的题目，写一篇题解记录一下。

HDU-7458 旅行(on Vjudge)

题面

题目描述

有一棵 \(n\) 个结点的无根树，每个结点都有对应的类型 \(c_i\) 和权重 \(w_i\) ，你需要在这棵树上规划若干次旅行。

对于一次旅行，你将从一个树上的一个结点出发，沿着树上的边进行旅行，最终到达另一个和起点类型相同的结点。

你会进行很多次旅行，但你希望对于每个结点，在所有旅行路线中最多只会经过一次。

一次旅行的价值是起始点和终止点的权重和，你需要规划旅行的方案使得旅行的总权重和最大。

输入格式

第一行输入一个正整数 \(t\ (1\le t\le 3)\) 代表数据组数。

对于每组数据：

第一行输入一个正整数 \(n\ (1\le n \le 2\times 10^5)\)，代表树的点数。

第二行输入 \(n\) 个正整数 \(c_i\ (1\le c_i \le n)\) ，代表每个结点的类型。

第三行输入 \(n\) 个正整数 \(w_i\ (1\le w_i \le n)\) ，代表每个结点的权重。

接下来 \(n-1\) 行每行两个正整数 \(u_i,v_i\ (1\le u_i,v_i\le n,u_i\neq v_i)\) ,代表树上一条边，输入保证是一棵树。

输出格式

输出一行一个正整数代表答案。

输入输出样例

输入 #1

1
7
3 1 1 2 2 2 3
2 4 1 5 4 6 2
1 2
1 3
2 4 
2 5
3 6 
3 7

输出 #1

13

Hint

一种最优的旅行方案为：

旅行路线1：\(4\to 2\to 5\) ，价值为 \(9\)。

旅行路线2：\(1\to 3\to 7\)，价值为 \(4\)。

价值总和为 \(13\)。

题解

我把本题的限制看成了恰好每个点经过一遍，于是浪费了一个晚上，在此警示后人。

注意到如果根节点走过或没走过状态确定，可以划分到每棵子树，考虑是树形 DP。设状态 \(f_{x,0}\) 表示子树 \(x\) 内 \(x\) 未走过的最大权值， \(f_{x,1}\) 表示子树 \(x\) 内 \(x\) 已走过的最大权值。\(f_{x,0}\) 的转移时容易的，因为不可能出现跨子树合并的贡献。

\[f_{x,0}=\sum_{v\in \operatorname{son}(x)}\max\{f_{x,0},f_{x,1}\} \]

接下来我们考虑 \(f_{x,1}\)，显然我们需要在子树 \(x\) 内寻找一个与 \(x\) 类型相同的点，然后合并。或者找到两个不同子树类型相同的点，合并顺便经过 \(x\)。考虑选择 \(y\) 作为合并的节点之一的贡献。设 \(x\to y\) 路径上的点为 \(a_1=x,a_2,\dots a_{k-1},a_k=y\)，则可以算出贡献为下面式子。

\[w_x+w_y+f_{y,0}+\sum_{i=1}^{k-1}f_{a_i,0}-\max\{f_{a_{i+1},0},f_{a_{i+1},1}\} \]

这个式子没有对齐，并不方便计算，于是我们稍微修改一下求和顺序，把下标对齐。

\[w_x+f_{x,0}+w_y+\sum_{i=2}^{k}f_{a_i,0}-\max\{f_{a_{i},0},f_{a_{i},1}\} \]

选择两颗颜色相同的子树合并，就是 \(x\) 到这两个点的路径的贡献减去 \(f_{x,0}\)，因为每一棵子树都恰好被按照 \(f_{x,0}\) 的方式多算了一次，直接减掉即可。代码中化简了这个式子。

注意到 \(w_x+f_{x,0}\) 是确定的，对于某种颜色的转移其实只需要使 \(w_y+\sum_{i=2}^{k}f_{a_i,0}-\max\{f_{a_{i},0},f_{a_{i},1}\}\) 最大。我们不妨把某棵子树某种颜色的最大值的集合存进一个 map。注意到这个式子只与 \(x\to y\) 的路径有关，在搜索回溯的时候加入当前节点，并对这个子树已有状态的集合使用懒标记增加逐步自底向上求出。我们还要支持合并两个集合，因为我们要合并多个子树的信息。

注意到这个合并可以使用启发式合并，对于每种颜色，选择最大的两个合并或选择根合并，顺便更新最大值的集合，就可以求出 \(f_{i,1}\)。时间复杂度 \(O(n\log^2 n)\)。

代码

顺便积累一个写法，C++17 以后，可以利用下列方式访问 map 中所有元素。(\([c,p]\) 中 \(c\) 为下标，\(p\) 为元素值)

for(auto [c,p]:g[b[v]])

#include <bits/stdc++.h>
using namespace std;
struct edge
{
	long long v,nxt;
}e[600000];
long long t,n,u,v,h[300000],b[300000],c[300000],w[300000],f[300000][2],mx[300000],tg[300000],cnt=0;
map<long long,long long>g[300000];
inline long long read()
{
	long long x=0,f=1;char ch=getchar();
	while (ch<'0'||ch>'9'){if (ch=='-') f=-1;ch=getchar();}
	while (ch>='0'&&ch<='9'){x=x*10+ch-48;ch=getchar();}
	return x*f;
}

void add_edge(long long u,long long v)
{
	e[++cnt].nxt=h[u];
	e[cnt].v=v;
	h[u]=cnt;
}

void dsu(long long x,long long fa)
{
	f[x][0]=f[x][1]=mx[x]=0;
    for(int i=h[x];i;i=e[i].nxt)
        if(e[i].v!=fa)dsu(e[i].v,x),f[x][0]+=mx[e[i].v];
    for(int i=h[x];i;i=e[i].nxt)
        if(e[i].v!=fa)
           {
           	long long v=e[i].v;
           	if(g[b[v]].count(c[x]))f[x][1]=max(f[x][1],g[b[v]][c[x]]+tg[b[v]]+f[x][0]+w[x]);
           	if(g[b[v]].size()>g[b[x]].size())swap(b[v],b[x]);
           	for(auto [c,p]:g[b[v]])
           	    if(g[b[x]].count(c))f[x][1]=max(f[x][1],p+tg[b[v]]+g[b[x]][c]+tg[b[x]]+f[x][0]);
           	for(auto [c,p]:g[b[v]])
           	    if(g[b[x]].count(c))g[b[x]][c]=max(g[b[x]][c],p+tg[b[v]]-tg[b[x]]);
           	    else g[b[x]][c]=p+tg[b[v]]-tg[b[x]];
		   }
	if(g[b[x]].count(c[x]))g[b[x]][c[x]]=max(g[b[x]][c[x]],w[x]-tg[b[x]]);
	else g[b[x]][c[x]]=w[x]-tg[b[x]];
    mx[x]=max(f[x][0],f[x][1]),tg[b[x]]+=f[x][0]-mx[x];
}

int main()
{
	t=read();
	while(t--)
	   {
	   	n=read();
	   	for(int i=1;i<=n;i++)h[i]=tg[i]=cnt=0,b[i]=i,g[i].clear();
	   	for(int i=1;i<=n;i++)c[i]=read();
	   	for(int i=1;i<=n;i++)w[i]=read();
	   	for(int i=1;i<=n-1;i++)u=read(),v=read(),add_edge(u,v),add_edge(v,u);
	   	dsu(1,0);
	   	printf("%lld\n",mx[1]);
	   }
	return 0;
}