Educational Codeforces Round 2 E – Lomsat gelral
链接:https://codeforces.com/contest/600/problem/E
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
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.
Leave A Comment