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

小美拿到了一棵树，每个节点有一个权值。初始每个节点都是白色。
小美有若干次操作，每次操作可以选择两个相邻的节点，如果它们都是白色且权值的乘积是完全平方数，小美就可以把这两个节点同时染红。
小美想知道，自己最多可以染红多少个节点？

public class Main {
    //题目给出n个点和n-1条边的一棵树 可以当成有向无环图来存
    //建树完之后从叶子节点一直往上贪心
    //从树的叶子开始涂色，才是可以涂最多的方法。
    static final int N = (int) (1e5 + 10);
    static int n;
    static long w[] = new long[N]; //结点权重数组
 
    static ArrayList<Integer> g[] = new ArrayList[N]; //邻接表
 
    static boolean st[] = new boolean[N]; //标记是否染色
    static int ans;//结果
 
    private static boolean fun(long x) {
        long d = (long) Math.sqrt(x);
        return d * d == x;
    }
 
    private static void dfs(int u, int fa) {
        for (Integer x : g[u]) {
            if (x == fa) continue; //确保不访问父节点
            dfs(x, u);
            if (fun(w[u] * w[x]) && (!st[x] && !st[u])) {
                ans += 2;
                st[x] = true;
                st[u] = true;
            }
        }
    }
 
    public static void main(String[] args) throws Exception {
        Scanner sc = new Scanner(System.in);
        n = sc.nextInt();
        for (int i = 0; i < N; i++) {
            g[i]=new ArrayList<>();
        }
        for (int i = 1; i <= n; i++) {
            w[i] = sc.nextInt();
        }
        for (int i = 0; i < n - 1; i++) {
            int u = sc.nextInt(), v = sc.nextInt();
            g[u].add(v);
            g[v].add(u);
        }
        dfs(1, -1); //从第一个结点开始dfs,其没有父结点设为-1
        System.out.println(ans);
    }
 
    //ps:由于只在乎相邻节点权值乘积为完全平方数的，所以不必使用所有的边建立树，而是使用乘积为平方数的边构造一棵树，然后对这棵树进行深度优先搜索
 
}