树上启发式合并详解 + 模板 CodeForces 600E Lomsat gelral —— JAVA

传送门：CodeForces 600E Lomsat gelral

直接先上题目然后讲吧。

E. Lomsat gelral

Time Limit: 2 Sec Memory Limit: 256 MB

You are given a rooted tree with root in vertex 1. Each vertex is coloured in some colour.

Let’s call colour c dominating in the subtree of vertex v if there are no other colours that appear in the subtree of vertex v more times than colour c. So it’s possible that two or more colours will be dominating in the subtree of some vertex.

The subtree of vertex v is the vertex v and all other vertices that contains vertex v in each path to the root.

For each vertex v find the sum of all dominating colours in the subtree of vertex v.

Input
The first line contains integer n (1 ≤ n ≤ 105) — the number of vertices in the tree.

The second line contains n integers ci (1 ≤ ci ≤ n), ci — the colour of the i-th vertex.

Each of the next n - 1 lines contains two integers xj, yj (1 ≤ xj, yj ≤ n) — the edge of the tree. The first vertex is the root of the tree.

Output
Print n integers — the sums of dominating colours for each vertex.

Examples

input：
4
1 2 3 4
1 2
2 3
2 4
output：
10 9 3 4

input：
15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13
output：
6 5 4 3 2 3 3 1 1 3 2 2 1 2 3

题目大意：

给你一颗n个结点的树，每个结点有一个数字代表颜色，随后n-1行为多条无向边，根节点必定是1号节点。输出的是，对于每个结点的子结点中，包含的颜色种类做多的那个颜色，若多个颜色是最多的，那么输出这些颜色的总和。

思路：

一般的暴力做法有：我给每一个结点做一个数组，然后把递归把的子结点的数字值加起来就是当前结点的颜色总量,取最多的就行了，时间复杂度是O(n)，但是空间复杂度爆炸大！还有一种就是，只建一个数组，对于每个结点都跑一遍他的子结点，每次跑完还需要清空数组，以便算下一个结点用。这样空间复杂度是下来了，变成O(n)，比较合理，但是时间复杂度直接变成O(n²)。

显然我们需要优化第二种方法的时间，而树上启发式搜索就是优化了这个暴力解法。

首先了解两个概念：

数组	解释
son[u]	保存u的重节点
size[u]	保存以u为根的所有子节点数量

重儿子：父结点的所有儿子节点中size[]最大的节点。
轻儿子：除重儿子以外的所有节点。

如上图，我遍历到了结点5，而4,6,9,10的值已经求完，我接下来需要求结点2，我们会发现，结点5做完后的数组并不需要去清空，只需要再dfs一遍其他结点与其求和，就能求出结点2的值。也就是说，我可以保留最后一个子结点的值，再拿它和其他子结点求和，这样就节省了一个子结点的dfs时间。

那么怎么做到收益最大化呢？显然我只需要让最后一个遍历的子结点为重儿子就可以了，经过大佬的分析，时间复杂度最多是nlogn，挺优秀的。

接下来是含注释的AC代码：

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.StreamTokenizer;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;
import java.util.Stack;
import java.util.TreeSet;



public class Main
{
	static int sz[],size[],son[],max,cnt[],o,p;
	static long ans[],sum[];
	static int top[],next[],end[];//邻接表数组
	static void con(int a,int b)//经典邻接表构建
	{
		next[++o]=top[a];
		top[a]=o;
		end[o]=b;
	}
	static void dfs(int x,int fu)//求重儿子
	{
		size[x]=1;
		for (int i=top[x];i!=0;i=next[i])
		{
			if (end[i]!=fu)//如果不是父结点
			{
				dfs(end[i],x);
				size[x]+=size[end[i]];
				if (size[son[x]]<size[end[i]])son[x]=end[i];//如果有子结点数更大的的结点，更新重儿子
			}
		}
	}
	static void g(int x,int fu,int s)
	{
		sum[cnt[sz[x]]]-=sz[x];//cnt[sz[x]]默认是0的，我们不使用索引号0的位置，所以0位置可以为负数
		cnt[sz[x]]+=s;//s用来做更改该结点是加出现次数还是减出现次数，就是+1和-1
		sum[cnt[sz[x]]]+=sz[x];//更改后的数量位置增加该值
		if (sum[max+1]>0)max++;//更新次数最大值
		if (sum[max]==0)max--;
		//上面这几句话个人觉得是最难理解的，也很妙，博客最后再单独说一下吧
		for (int i=top[x];i!=0;i=next[i])
		{
			if (end[i]!=p && end[i]!=fu)//因为重节点的值是保留的，也就是不需要dfs的，所以end[i]!=p避免重结点的值重复操作
			{
				g(end[i],x,s);
			}
		}
		
	}
	static void dfs1(int x,int fu,int s)//当前结点，父结点，s==1代表当前是重结点，需要保留
	{
		for (int i=top[x];i!=0;i=next[i])
		{
			if (end[i]!=son[x] && end[i]!=fu)//是重儿子或是父结点都不递归
			{
				dfs1(end[i],x,0);
			}
		}
		if (son[x]!=0)//如果不是叶子结点，就有重儿子，那么最后这里再递归重儿子
		{
			dfs1(son[x],x,1);
			p=son[x];//标记重结点，免得dfs求值重结点重复记录
		}
		g(x,fu,1);//dfs一遍该结点的子结点，添加数组值
		ans[x]=sum[max];//记录答案
		p=0;//当前结点求完了，可能需要dfs一遍清空数组，不能保留该重结点的值了
		if (s==0)g(x,fu,-1);//如果不是重结点，删除数组，相比起for一遍整个数组清空速度快很多
	}
	public static void main(String[] args) throws IOException
	{
		int n=ini();
		p=0;//标记重儿子
		o=0;//邻接表用来做索引的值
		sz=new int [n+1];//每个结点的颜色数字
		ans=new long [n+1];//答案
		
		//三个邻接表数组
		top=new int [n+1];
		next=new int [n<<1];
		end=new int [n<<1];
		
		size=new int [n+1];//记录子结点数量
		son=new int [n+1];//记录重儿子
		sum=new long [n+1];//记录该次数的颜色总和
		cnt=new int [n+1];//记录该颜色的出现次数
		for (int i=1;i<=n;i++)sz[i]=ini();
		for (int i=1;i<n;i++)
		{
			int a=ini();
			int b=ini();
			con(a,b);
			con(b,a);
		}
		dfs(1,0);
		dfs1(1,0,1);
		for (int i=1;i<=n;i++)
		{
			if (i==1)out.print(ans[i]);
			else out.print(" "+ans[i]);
		}
		
		out.flush();
	}
/*


*/
	static StreamTokenizer in=new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
	static PrintWriter out=new PrintWriter(new BufferedWriter(new OutputStreamWriter(System.out)));
	static int ini() throws IOException
	{
		in.nextToken();
		return (int)in.nval;
	}
	static long inl() throws IOException
	{
		in.nextToken();
		return (long)in.nval;
	}
	static String ins() throws IOException
	{
		in.nextToken();
		return in.sval;
	}
	
}

针对上面51行开始的几句代码说一下，cnt[]记录的是每个颜色的出现次数，而sum[]是该出现次数的颜色总和，我当前这个颜色出现了cnt[sz[x]]次，但是现在我该颜色的出现次数要开始增加或者减少了，所以我需要将该次数的值减去该颜色的值，然后调整cnt[sz[x]]的颜色次数，在另一个次数上增加该颜色值。
最后两句if好理解，max就是记录最多的相同颜色是多少，通过询问max+1次是否有值与max次的值是否一个颜色都没了来更新。

	sum[cnt[sz[x]]]-=sz[x];
	cnt[sz[x]]+=s
	sum[cnt[sz[x]]]+=sz[x];
	if (sum[max+1]>0)max++;
	if (sum[max]==0)max--;