题解 | #蚂蚁森林之王#

可以将这题转化为:计算森林中每棵树的每个子树的结点个数.

当然无需真的去建树,只需要用一个二维数组来模拟树结构即可(记录每个节点的孩子)

转换规则: 若a崇拜b,即a是b的孩子.

import java.util.*;
public class Main{
    public static void main(String[]args){
        new Main().solve();
    }
    int n;
  
    int[]worship;
  
    List<List<Integer>>tree;
  
    long[]ans;
  
    void solve(){
        Scanner sc=new Scanner(System.in);
        n=sc.nextInt();
        worship=new int[n+1];
        for(int i=0;i<n;i++){
            worship[i+1]=sc.nextInt();
        }
        tree=new ArrayList<>(n+1);
        tree.add(null);
        for(int i=1;i<=n;i++){
            tree.add(new ArrayList<>());
        }
        for(int i=2;i<=n;i++){
            if(worship[i]!=0)
                tree.get(worship[i]).add(i);	//i是worship[i]的孩子
        }
        ans=new long[n+1];
        Arrays.fill(ans,-1);
        for(int i=1;i<=n;i++)
            dfs(i);
        for(int i=1;i<=n;i++)
            System.out.println(ans[i]);
    }
    //递归计算子树所含的节点个数
    long dfs(int k){
        if(ans[k]!=-1)	//避免重复计算
            return ans[k];
        List<Integer>children=tree.get(k);
        long sum=1;
        for(int i=0;i<children.size();i++){
            sum+=dfs(children.get(i));
        }
        ans[k]=sum;
        return sum;
    }
}