百度笔试-java-第三道红蓝树

笔试的时候没调出来，结束了之后debug修改的。自己测试了几个用例没问题，不知道能不能全部通过。

import java.util.*;

public class Main {

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        char[] c = sc.next().toCharArray();
        TreeNode[] nodes = new TreeNode[n];
        for (int i = 0; i < n; i++) {
            nodes[i] = new TreeNode(c[i]);
        }
        for (int i = 0; i < (n - 1); i++) {
            int u = sc.nextInt()-1;
            int v = sc.nextInt()-1;
            nodes[u].treeNodes.add(nodes[v]);
            nodes[v].treeNodes.add(nodes[u]);
        }
        TreeNode root = nodes[0];
        System.out.println(getSum(root)/2);
    }

    static Set<TreeNode> getSumSet = new HashSet<>();

    static int getSum(TreeNode root) {
        if (root == null || getSumSet.contains(root)) return 0;
        getSumSet.add(root);
        int sum = 0;
        for (TreeNode node : root.treeNodes) {
            int rootCount = getCount(root, node, 1);
            int nodeCount = getCount(node, root, 1);
            sum += Math.abs(rootCount - nodeCount);
        }
        for (TreeNode node : root.treeNodes) {
            sum += getSum(node);
        }
        return sum;
    }


    static int getCount(TreeNode root, TreeNode dis, int sum) {
        if (root == null) return 0;
        for (TreeNode node : root.treeNodes) {
            if (node == dis) continue;
            sum += (root.color == node.color ? 0 : 1);
        }
        for (TreeNode node : root.treeNodes) {
            if (node == dis) continue;
            sum += getCount(node, root, 0);
        }
        return sum;
    }
}

class TreeNode {
    public char color;
    public List<TreeNode> treeNodes;

    public TreeNode(char color) {
        this.color = color;
        treeNodes = new ArrayList<>();
    }
}

用例1：

输入：
4
BBRR
1 2
3 2
4 1
输出：2

用例2：

输入：
5
BBRRR
1 2
3 2
4 1
2 5
输出：7

用例3：

输入：
7
BBRRRBBR
1 2
3 2
4 1
2 5
4 6
7 5
输出：17