传送门:CodeForces 600E Lomsat gelral
直接先上题目然后讲吧。
E. Lomsat gelral
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--;