题面

核心思想

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

代码

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)
                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]的最大值
                // 用当前节点的dp[0]减去就剩下了其他孩子的最大值和
                // 那么当前节点和孩子节点nxt染色 dp[1]就等于 1. nxt这个孩子节点的dp[0] + 2. 其他孩子节点的dp最大值的和
                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;
    }
}