MC0483过园数统计

大观园内众姐妹成立了诗社。各社员居所相连正如一棵树，居所的编号为1∼n1∼n。各社员均有一个编号，或为诗社核心成员（记为 1），或为列席成员（记为 0）。

若以某姐妹的居所为 “社集之地”（即视其为根），自社集之地至另一居所须过几处园门，便称为那居所的 “过园数”（即根到该点的边数）。

现在做为社长的小码妹需要知道：对于园中每处居所i(1≤i≤n)i(1≤i≤n)，若以该点为社集之地时，所有诗社核心成员（即记为1者）的居所，其 “过园数” 共有多少种不同的数值？

格式

输入格式：

第一行一个整数n(1≤n≤3×104)n(1≤n≤3×104)。
第二行nn个整数a1∼an(0≤ai≤1)a1​∼an​(0≤ai​≤1)。
接下来n−1n−1行，每行两个整数u,v(1≤u,v≤n)u,v(1≤u,v≤n)，表示点uu和点vv之间有一条边。

输出格式：

输出nn行，其中第ii行代表当点ii为根节点时的答案。

样例 1

输入：

3 0 1 1 1 2 1 3

复制

输出：

1 2 2

复制

样例 2

输入：

6 0 0 1 1 1 1 1 2 1 3 1 4 2 5 3 6

复制

输出：

2 3 4 3 3 4

复制

备注

样例11解释：
当11号点为根节点时，值为11的点（22号点和33号点）到达11号点的距离分别为1,11,1，一共有11种不同的数值。
当22号点为根节点时，值为11的点（22号点和33号点）到达22号点的距离分别为0,20,2，一共有22种不同的数值。
当33号点为根节点时，值为11的点（22号点和33号点）到达33号点的距离分别为2,02,0，一共有22种不同的数值。

代码1（超时）

import java.io.BufferedReader; import java.io.BufferedWriter; import java.io.IOException; import java.io.InputStreamReader; import java.io.OutputStreamWriter; import java.util.*; public class Main { static int N=3*10010; static int a[]=new int[N]; static boolean f[]=new boolean[N]; static List<Integer>[] adj=new List[N]; public static void main(String[] args) throws IOException { BufferedReader br = new BufferedReader(new InputStreamReader(System.in)); BufferedWriter bw=new BufferedWriter(new OutputStreamWriter(System.out)); StringTokenizer st=new StringTokenizer(br.readLine()); int n=Integer.parseInt(st.nextToken()); st=new StringTokenizer(br.readLine()); for (int i = 1; i <=n; i++) { a[i]=Integer.parseInt(st.nextToken()); adj[i]=new ArrayList<>(); } for (int i = 1; i < n; i++) { st=new StringTokenizer(br.readLine()); int a=Integer.parseInt(st.nextToken()),b=Integer.parseInt(st.nextToken()); adj[a].add(b); adj[b].add(a); } for (int i = 1; i <= n; i++) { for (int j = 1; j <= n; j++) { f[j]=false; } int dis[]=new int[n+1]; dis[i]=0; Set<Integer> set=new HashSet<>();//有多少不同的边数 if(a[i]==1)set.add(0); Queue<Integer> queue=new LinkedList<>(); queue.add(i);f[i]=true; while(!queue.isEmpty()){ int top=queue.poll(); for(int son:adj[top]){ if(f[son]==true)continue;//相当于是根节点跳过 dis[son]=dis[top]+1; if(a[son]==1)set.add(dis[son]); queue.add(son); f[son]=true; } } System.out.println(set.size()); } br.close(); bw.flush(); bw.close(); } }

代码2

import java.io.BufferedReader; import java.io.BufferedWriter; import java.io.IOException; import java.io.InputStreamReader; import java.io.OutputStreamWriter; import java.util.*; public class Main { static int N=3*10010,n,words;//words表示需要多少的long块来记录 static int a[]=new int[N]; static int ans[]=new int[N]; static long f[][];//表示以i为根节点 所有核心店到i的距离 //比如 f[u] 中 第一个long 表示距离0-63 是否存在 存在则是1 不存在则是0 0..1010 表示距离1、2存在 //第二个long 表示距离64-127 是否存在 存在则是1 不存在则是0 0..1001 表示距离 64、67存在 static List<Integer>[] adj=new List[N];//邻接表 public static void main(String[] args) throws IOException { BufferedReader br = new BufferedReader(new InputStreamReader(System.in)); BufferedWriter bw=new BufferedWriter(new OutputStreamWriter(System.out)); StringTokenizer st=new StringTokenizer(br.readLine()); n=Integer.parseInt(st.nextToken()); words=(n+63)/64; f=new long[N][words]; st=new StringTokenizer(br.readLine()); for (int i = 1; i <=n; i++) { a[i]=Integer.parseInt(st.nextToken()); adj[i]=new ArrayList<>(); } for (int i = 1; i < n; i++) { st=new StringTokenizer(br.readLine()); int a=Integer.parseInt(st.nextToken()),b=Integer.parseInt(st.nextToken()); adj[a].add(b); adj[b].add(a); } dfs1(1,0);//以 1 为根的树形下，f[u] 的含义就是“u 子树内的所有核心点到 u 的距离集合”。 dfs2(1,0,new long[words]);//不断的换根 求过圆数 开始时 1的外部节点 也就是向上的点集合到1为空集 StringBuilder stringBuilder=new StringBuilder(); for (int i = 1; i <= n; i++) { stringBuilder.append(ans[i]); stringBuilder.append("\n"); } bw.write(stringBuilder.toString().trim()); br.close(); bw.flush(); bw.close(); } static void dfs1(int u,int p){//计算以每个节点为根 该子树中核心点到根节点的边数 if(a[u]==1){//自己到自己 距离为0 setBit(f[u],0); } for(int v:adj[u]){ if(v==p)continue; dfs1(v,u); or(f[u],shiftLeft(f[v],1));//合并 } } static void dfs2(int u,int p,long upSet[]){//不断的换根 求出最终结果 //upSet表示根节点向上的外面的核心点距离集合 long all[]=Arrays.copyOfRange(f[u], 0,words); or(all,upSet); for (int i = 0; i < words; i++) { ans[u]+=Long.bitCount(all[i]);//bigCount统计1的个数 } //以下全部都是为换根做准备的 // 比如u的三个叶子结点是 v1 v2 v3 v4 v5 //当v3作为根的时候 我们需要求出v1-3 的所有或的距离的结果 以及v4-5 的所有或的距离的结果 //如果每次都把前面和后面的距离全部或起来 复杂度很高 //我们可以预处理出前缀或 以及 后缀或 //统计孩子结点 这里要有索引区分子树 因为后续要用索引计算前缀 后缀或 int k=0; List<Integer> list=new ArrayList<>(); for(int v:adj[u]){ if(v==p)continue; k++; list.add(v); } if(k==0)return; //计算每个孩子节点（子树）到u的距离 long g[][]=new long[k][words];//表示第i课子树的核心点到u的距离的集合 for (int i = 0; i < k; i++) { g[i]=shiftLeft(f[list.get(i)], 1); } long prefix[][]=new long[k][words];//前缀或 prefix[0]=g[0]; for (int i = 1; i < k; i++) { or(prefix[i],prefix[i-1]); or(prefix[i],g[i]); } long suffi[][]=new long[k][words];//后缀或 suffi[k-1]=g[k-1]; for (int i = k-2; i>=0; i--) { or(suffi[i], suffi[i+1]); or(suffi[i], g[i]); } //依次把每个孩子节点当成根 //当一个孩子结点变为根的时候 它的upSet集合就变成了 我这个孩子结点之前的所有结点 还有u //还有我这个孩子结点之后的所有结点 还有传进来的upSet for (int i = 0; i < k; i++) { int v=list.get(i); //v的upSet先计算成到u的距离 最后统一加1 long other[]=new long[words]; if(a[u]==1)setBit(other, 0); or(other,upSet); if(i-1>=0)or(other, prefix[i-1]); if(i+1<k)or(other, suffi[i+1]); long newUp[]= shiftLeft(other, 1); dfs2(v,u,newUp); } } static void or(long fina[],long src[]){ for (int i = 0; i < words; i++) { fina[i]|=src[i]; } } static long[] shiftLeft(long t[],int x){//将t中所有的距离加上x long k[]=new long[words]; if((x&63)==0){ int change = (x>>6); for (int i = 0; i < words; i++) { if(i+change<words)k[i+change]=t[i]; } }else{ for (int i = 0; i < words; i++) { long w=t[i]; int id=i+(x>>6); if(id<words)k[id]|= w<<(x&63); if(id+1<words)k[id+1]|= w>>>(64-(x&63)); // >> 是算术右移，如果 w 的最高位是 1，右移后左边会补 1（符号位），导致结果错误。 // 位图操作要求逻辑右移，即左边一律补 0，必须用 >>>。 } } return k; } static void setBit(long t[],int d){//将数组表示的d的距离设置为1 t[d>>6]|= (1L<<(d & 63));//t[d/64]|= (1<<(d % 64)) #必须用 1L 1 是 int 类型，只有 32 位 //>> 6 相当于除以 64（因为 2⁶ = 64） d >> 6 是找到距离 d 应该放在哪个 long 里 //& 63 相当于模 64 } }