题目链接:Cut the tree
题解:题目要求求一条边,去掉这条边后得到的两棵树的节点和差的绝对值最小。
暴力求解会超时。
如果我们可以求出以每个节点为根的子树的节点之和,那么当我们去掉一条边(a,b)的时候,其中的一棵树必是以a或者b为根的子树,那么我们就可以知道生成的两棵树的节点之和了。所以,当我们得到以每个节点为根的子树节点和这个信息后(把这个信息存储在TreeNode的value变量中),通过遍历每个节点(每个节点为根的子树必然为某次cut得到的子树之一,例如下图中,2为根的子树对应切割边(1,2)得到的树,而5为根的子树对应切割边(1,5)得到的树),求出最小的W-2*root.value即为所求(W是所有节点的值得和)。
树节点的数据结构定义如下:
static class TreeNode{ int value; int number; ArrayList<TreeNode> children; public TreeNode(int value,int number){ this.value = value; this.number = number; children = new ArrayList<TreeNode>(); } }
value表示该顶点对应的值,number表示顶点的编号,children表示子节点的列表(这里指表示连接关系,也有可能children里面其实存放了父节点,不过可以用接下来的visited数组避免访问到父节点)。
然后从树的任意一个节点开始dfs,把它的value值依次加上各个child为root的子树节点值,就得到以它为root的子树的节点值得和了。而它的child为root的子树节点和就用递归的方法求。上面说了,每次读入一条边(a,b)的时候,即把b加入到a的children集合里面,也把a加入到b的children集合里面。那么在dfs的时候,怎么知道谁是父节点,谁是子节点呢?就用一个visited数组,因为在dfs过程中,父节点一定先被访问,所以在一个节点的children集合里面,已经被visited过的节点就不是它的子节点了。
最后代码如下:
1 import java.util.*; 2 3 public class Solution { 4 static class TreeNode{ 5 int value; 6 int number; 7 ArrayList<TreeNode> children; 8 public TreeNode(int value,int number){ 9 this.value = value; 10 this.number = number; 11 children = new ArrayList<TreeNode>(); 12 } 13 } 14 15 private static int dfs(TreeNode root,boolean[] visited){ 16 visited[root.number]=true; 17 for(TreeNode child:root.children){ 18 if(!visited[child.number]){ 19 root.value += dfs(child, visited); 20 } 21 } 22 return root.value; 23 } 24 25 public static void main(String[] args) { 26 Scanner in = new Scanner(System.in); 27 int n = in.nextInt(); 28 TreeNode[] treeNodes = new TreeNode[n]; 29 int W = 0; 30 31 for(int i = 0;i < n;i++) 32 { 33 int value = in.nextInt(); 34 treeNodes[i]=new TreeNode(value,i); 35 W += value; 36 } 37 38 for(int i = 0;i < n-1;i++){ 39 int a = in.nextInt(); 40 int b = in.nextInt(); 41 treeNodes[a-1].children.add(treeNodes[b-1]); 42 treeNodes[b-1].children.add(treeNodes[a-1]); 43 } 44 45 boolean[] visited = new boolean[n]; 46 47 dfs(treeNodes[0], visited); 48 int miniDif = Integer.MAX_VALUE; 49 50 for(TreeNode t:treeNodes){ 51 miniDif = Math.min(miniDif, Math.abs(W-2*t.value)); 52 } 53 54 System.out.println(miniDif); 55 56 57 58 } 59 }
另外要积累的一点是java中nested class和inner class的区别,见stackoverflow上面的详细解答,所以上述代码中第4行要把TreeNode声明为nested class,否则它只能通过类Solution调用。