携程笔试8.30

import java.util.*;

public class Main {

    public static int n;
    public static int[] pre;
    public static int[] total = new int[3];
    public static List<Integer>[] tree;
    public static Map<Integer, int[]>[] cntMap;
    public static String s;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        n = sc.nextInt();
        pre = new int[n];
        tree = new List[n];
        for (int i = 0; i < n; i++) {
            tree[i] = new ArrayList<>();
        }
        cntMap = new Map[n];
        for (int i = 0; i < n; i++) {
            cntMap[i] = new HashMap<>();
        }

        sc.nextLine();
        s = sc.nextLine();
        for (int i = 0; i < s.length(); i++) {
            char c = s.charAt(i);
            calc(total, c);
        }

        for (int i = 0; i < n-1; i++) {
            int u = sc.nextInt();
            int v = sc.nextInt();
            tree[u-1].add(v-1);
            tree[v-1].add(u-1);
            pre[u-1]++;
            pre[v-1]++;
        }

        initCntMap();
        System.out.println(solution());
    }

    private static int solution() {
        int res = 0;
        for (int i=0; i<n; ++i) {
            for (Map.Entry<Integer, int[]> entry : cntMap[i].entrySet()) {
                int[] cnt = entry.getValue();
                if (judge(cnt))
                    res++;
            }
        }
        return res;
    }

    private static void initCntMap() {
        Queue<Integer> q = new LinkedList<>();
        for (int i=0; i<n; ++i) {
            if (pre[i] == 1)
                q.offer(i);
        }
        while (!q.isEmpty()) {
            int u = q.poll();
            int[] cnt = new int[3];
            calc(cnt, s.charAt(u));
            int k = -1;
            for (int i = 0; i < tree[u].size(); i++) {
                int v = tree[u].get(i);
                pre[v]--;
                if (pre[v] == 1)
                    q.offer(v);
                if (cntMap[v].containsKey(u)) {
                    merge(cnt, cntMap[v].get(u));
                } else {
                    k = v;
                }
            }
            if (k >= 0)
                cntMap[u].put(k, cnt);
        }
    }

    private static void merge(int[] target, int[] other) {
        for (int i=0; i<3; ++i)
            target[i] += other[i];
    }

    private static void calc(int[] cnt, char c) {
        if (c == 'r')
            cnt[0]++;
        else if (c == 'g')
            cnt[1]++;
        else if (c == 'b')
            cnt[2]++;
    }

    private static boolean judge(int[] cnt) {
        for (int i=0; i<3; ++i) {
            if (cnt[i] == 0 || cnt[i] == total[i])
                return false;
        }
        return true;
    }
}