小美的树上染色（美团2024届秋招笔试第一场编程真题）

题面

核心思想

树形DP
dp[1]表示以当前节点为根节点以及所包含的子树 且 当前节点能染色的最大染色数量
dp[0]表示以当前节点为根节点以及所包含的子树 且 当前节点不染色的最大染色数量

dp[0]好算，只用把孩子节点的dp值取max累加就好了

dp[1] 假设当前节点为cur，cur能和孩子节点nxt染色

那么dp[1] = sum(max(dp[other][0], dp[other][1]))（除了nxt的其他孩子节点的dp累加值）+ dp[nxt][0]（nxt节点的dp[0]） + 2（cur和nxt染色产生的贡献）

详情看注释~

代码

import java.util.*;

public class Main {
    public static void main(String[] args) {
        final long MOD = (long) (1e9 + 7);
        Scanner scanner = new Scanner(System.in);

        int n = scanner.nextInt();
        long[] value = new long[n + 1];
        List<Integer>[] next = new List[n + 1];
        //存放value
        for(int i = 1; i <= n; i++){
            value[i] = scanner.nextInt();
            next[i] = new ArrayList<>();
        }
        //建树
        for(int i = 1; i < n; i++){
            int x = scanner.nextInt();
            int y = scanner.nextInt();
            next[x].add(y);
            next[y].add(x);
        }
        int[] res = dpOnTheTree(1, -1, value, next);
        System.out.println(Math.max(res[0], res[1]));
    }

    //dp[0] 表示当前节点为根 当前节点不染色的最大染色数量  dp[1]则表示当前节点染色的最大染色数量
    static int[] dpOnTheTree(int cur, int pre, long[] value, List<Integer>[] next){
        int[] dp = new int[2];
        //存放孩子节点的dp结果
        HashMap<Integer, int[]> res = new HashMap<>();

        //当前节点的dp结果分两步做
        //dp[0]
        for(int nxt: next[cur]){
            if(nxt == pre) //pre为父节点 防止走回去咯~
                continue;
            int[] child = dpOnTheTree(nxt, cur, value, next);
            res.put(nxt, child);

            // 当前节点不染色 那就是所有孩子节点的最大值和
            dp[0] += Math.max(child[0], child[1]);
        }

        //dp[1]
        for(int nxt: next[cur]){
            if(nxt == pre) 
                continue;
            long mul = value[cur] * value[nxt];
            long sqrt = (long) Math.sqrt(mul);
            // 可以和孩子节点染色
            if(sqrt * sqrt == mul){
                // dp[0] 存放的是所有孩子节点染色或不然染色的最大值和
                // Math.max(res.get(nxt)[0], res.get(nxt)[1]) 取需要染色的孩子节点nxt的dp[1], dp[0]的最大值
                // 以上两步计算结果 相减 就剩下了其他孩子的最大值和
                // 那么当前节点和孩子节点nxt染色 dp[1] = 其他孩子节点的dp最大值的和（也就是前两步相减的结果）+ nxt这个孩子节点的dp[0] + 2 
                dp[1] = Math.max(dp[1], dp[0] - Math.max(res.get(nxt)[0], res.get(nxt)[1]) + res.get(nxt)[0] + 2);
            }
        }
        return dp;
    }
}