滴滴9.28笔试

第二题：

package whut.yy.cainiao;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

/**
 * @author gagayang
 * @description
 * @date 2023/9/28
 */
public class Main01 {
    // 3 3 <-- 总结点数、询问次数
    // 1 1 <-- 除去根节点1，记录其余节点的父节点
    // 1 2 3 <-- 询问节点
    // return 2 3 3
    public static void main(String[] args) {
        // 处理输入数据
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int m = in.nextInt();
        int[] par = new int[n];
        for (int i = 0; i < n - 1; i++) {
            par[i + 1] = in.nextInt() - 1;
        }

        // 获取结果
        Solution solution = new Solution();
        int[] ans = solution.sumOfDistancesInTree(n, par);

        // 根据询问打印结果
        for (int i = 0; i < m; i++) {
            int node = in.nextInt() - 1;
            System.out.print(ans[node]);
            if (i < m - 1) {
                System.out.print(" ");
            }
        }
    }
}

class Solution {
    private List<Integer>[] g;
    private int[] ans;// 每个节点到其他节点的距离之和
    private int[] size;// 当前节点为根的子树的节点个数

    public int[] sumOfDistancesInTree(int n, int[] par) {
        g = new ArrayList[n]; // g[x] 表示 x 的所有邻居
        Arrays.setAll(g, e -> new ArrayList<>());
        // 根为0，不用设置父节点
        for (int i = 1; i < n; i++) {
            int parent = par[i];
            g[i].add(parent);
            g[parent].add(i);
        }
        ans = new int[n];
        size = new int[n];
        dfs(0, -1, 0); // 0 没有父节点
        reroot(0, -1); // 0 没有父节点
        return ans;
    }

    /**
     * 算 ans[0] + 算 size
     *
     * @param x
     * @param fa
     * @param depth
     */
    private void dfs(int x, int fa, int depth) {
        ans[0] += depth; // 先序位置算距离和： depth 为 0 到 x 的距离
        size[x] = 1;// 算子树的节点总数：根算一个，根的每个孩子节点构成的子树的节点个数之和也算
        for (int y : g[x]) { // 遍历 x 的邻居 y
            if (y != fa) { // 避免访问父节点
                dfs(y, x, depth + 1); // x 是 y 的父节点
                size[x] += size[y]; // 后序位置累加 x 的儿子 y 的子树大小
            }
        }
    }

    /**
     * 算 ans
     *
     * @param x
     * @param fa
     */
    private void reroot(int x, int fa) {
        for (int y : g[x]) { // 遍历 x 的邻居 y
            if (y != fa) { // 避免访问父节点
                ans[y] = ans[x] + g.length - 2 * size[y];
                reroot(y, x); // x 是 y 的父节点
            }
        }
    }
}