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
}
}