8.27 米哈游笔试

第三题就差一点写完，有没有大佬看一下对不对

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        Node[] list = new Node[n+1];
        in.nextLine();
        String s = in.nextLine();
        char[] c = s.toCharArray();
        for(int i = 1; i <= n; i++) {
            list[i] = new Node();
        }
        for(int i = 1; i <= n; i++) {
            list[i].c = c[i-1];
            if(i != n) {
                int a = in.nextInt();
                int b = in.nextInt();
                list[a].neighbor.add(list[b]);
                list[b].neighbor.add(list[a]);
            }
        }
        Node root = null;
        for(int i = 1; i <= n; i++) {
            if(list[i].neighbor.size() == 1) {
                root = list[i];
                break;
            }
        }
        root.num = dfs(root, null);//计算子树节点数量
        long sum = 0;
        for(int i = 1; i <= n; i++) {
            if(list[i].c == 'h') {
                long temp1 = 0;//"mhy"y一侧的节点数
                long temp2 = 0;//"mhy"m一侧的节点数
                for(Node node : list[i].neighbor) {
                    if(node.c == 'y') {
                        if(list[i].father == node) {
                            temp1 = n - node.num + 1;
                        } else {
                            temp1 = node.num;
                        }
                        for (Node node1 : list[i].neighbor) {
                            if(node1.c == 'm') {
                                if(list[i].father == node1) {
                                    temp2 = n - node1.num + 1;
                                } else {
                                    temp2 = node1.num;
                                }
                                sum += temp1 * temp2;
                            }
                        }
                    }
                }
            }
        }
        System.out.println(sum);
    }

    public static int dfs(Node node, Node father) {
        if(node == null) return 0;
        int sum = 1;
        node.father = father;
        for(Node n : node.neighbor) {
            if(n == father) continue;
            sum += dfs(n, node);
        }
        node.num = sum;
        return sum;
    }
}

class Node {
    char c;//字符
    List<Node> neighbor = new ArrayList<>();
    Node father;//父节点
    int num;//子树节点数量
}