Educational Codeforces Round 2 E – Lomsat gelral

链接:https://codeforces.com/contest/600/problem/E

public class TaskE {
    public void solve(int testNumber, InputReader in, OutputWriter out) {
        int n = in.readInt();
        int[] res = new int[n];
        int[] colorAry = in.readIntArray(n);
        //the number of occurences for each colour for each vertexs
        Map<Integer, Integer>[] colorsMap = new Map[n];
        //total count of sums for each vertexs(result)
        long[] totalColorSums = new long[n];
        // total count for each color
        int[] totalColorCounts = new int[n];
        //read graph
        ArrayList<Integer>[] g = new ArrayList[n];
        for(int i = 0 ; i < n ; i++) {
            g[i] = new ArrayList<>();
        }
        for(int i = 0 ; i < n - 1; i++){
            int a = in.readInt()-1;
            int b=  in.readInt()-1;
            g[a].add(b);
            g[b].add(a);
        }
        dfs(0, -1, colorAry, g, colorsMap,totalColorSums, totalColorCounts);
        out.printLine(totalColorSums);
    }

    private void dfs(int current, int parentNode, int[] colorAry,  ArrayList<Integer>[] g,Map<Integer, Integer>[] colorsMap,  long[] totalColorSums,  int[] totalColorCounts) {
        int currentColor = colorAry[current];
        Map<Integer, Integer> currentNodeColorCount = new HashMap<>();
        currentNodeColorCount.put(currentColor, 1);
        int count = 1;
        long sum = currentColor;
        for (int next : g[current]) {
            if (next == parentNode ) {
                continue;
            }
            //dfs to all connected nodes
            dfs(next, current, colorAry, g, colorsMap, totalColorSums, totalColorCounts);
            Map<Integer, Integer> connectedNodeColorCount = colorsMap[next];
            if (connectedNodeColorCount.size() > currentNodeColorCount.size()) {// merge small to large
                connectedNodeColorCount = currentNodeColorCount;
                currentNodeColorCount = colorsMap[next];// move to next
                count = totalColorCounts[next];
                sum = totalColorSums[next];
            }

            for (Map.Entry<Integer, Integer> connectedNode : connectedNodeColorCount.entrySet()) { //merge color and counts
                int connectedNodeColor = connectedNode.getKey();
                int connectedNodeCount = connectedNode.getValue();
                currentNodeColorCount.put(connectedNodeColor,  currentNodeColorCount.getOrDefault(connectedNodeColor, 0) +connectedNodeCount); // update current node's color count
                if (count < currentNodeColorCount.get(connectedNodeColor)) { // if total color is less than the current color count, in other words, we found larger count of some color.
                    count = currentNodeColorCount.get(connectedNodeColor);// then we need to update the current max color and count.
                    sum = connectedNodeColor;
                }
               else if (count == currentNodeColorCount.get(connectedNodeColor)){// if the total color is same to current color count, in other words,  two color has same color count and the count is the largest count.
                    sum += connectedNodeColor;// we need to add two color together.
                }
            }
        }

        colorsMap[current] = currentNodeColorCount;
        totalColorCounts[current] = count;
        totalColorSums[current] = sum;
    }
}

上面code的idea不是我自己想的, 是我看了答案后, 找到另一个人的解法, 然后优化出来的. 答案里写的方法有点太难了, 我看了另一个文章还有用dsu(union-find) in tree优化的. 题目边上的Problem tags里也有dsu. dsu在这里优化主要是用来减少dfs中合并子树中color count, 但是我感觉其实并不需要这样做:

首先, 用dsu就需要标记leaf和root, 然而我们的本意只是想找到leaf的color count, 然后合并, 所以其实只要记录并比较当前dfs节点上的count大小和我们记录下的与其相邻的节点大小, 就可以知道谁是leaf, 谁是root (dfs原理). 这样在dfs经过leaf的时候, 我们只需要记录当前的count就可以, 在dfs经过root的时候, 我们通过parent 这个变量, 就知道上次调用的节点数据.

其次, dsu合并的均摊是logn, 但是使用dsu需要维护另一个数组, 并且构图的时候也要多构造一个dsu. 而这里我们使用两个变量, sum和count, sum记录当前node下属所有节点的有最大count的color的合, count记录当前node下属所有节点的有最大count. 这里我们把这两个关键数据放在dfs的外侧, 作为闭包内的变量, 然后用totalColorCounts和totalColorSums记录全局数据. 这样我们在每次访问节点的时候, 就可以通过更新sum和count计算出当前节点的答案, 记录在
totalColorCounts和totalColorSums里.

这个题我觉得最难的是找到更新和维护connectedNodeColorCount和currentNodeColorCount, 一个是子节点的color和count数据, 另一个是这个子节点的父节点的color和count数据. 这两组数据决定了最后的total color和total sum. 另外, 语句connectedNodeColorCount.size() > currentNodeColorCount.size(), 确保了在当前节点为子节点的情况下, 需要交换connectedNodeColorCount和currentNodeColorCount, 这样就和dsu一样,确保了merge是从小的map到大的map.