Advertisement
Guest User

k distance

a guest
Feb 22nd, 2019
92
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Java 4.64 KB | None | 0 0
  1.     public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
  2.         ParentedTreeNode parentedTreeNode = constructParentedTree(root);
  3.  
  4.         List<Integer> distanceK = distanceK(parentedTreeNode, target.val, K);
  5.  
  6.         return distanceK;
  7.     }
  8.  
  9.     private List<Integer> distanceK(ParentedTreeNode treeNode, int val, int k) {
  10.         ParentedTreeNode target = FindTargetNode(treeNode, val);
  11.         List<Integer> distanceK = distanceK(target, k);
  12.         return distanceK;
  13.     }
  14.  
  15.     private List<Integer> distanceK(ParentedTreeNode target, int K) {
  16.         List<Integer> result = new ArrayList<>();
  17.         Queue<ParentedTreeNode> q = new LinkedList<>();
  18.         Queue<Integer> qk = new LinkedList<>();
  19.         q.add(target);
  20.         qk.add(K);
  21.         while (!q.isEmpty()) {
  22.             ParentedTreeNode node = q.remove();
  23.             int k = qk.remove();
  24.             if (k == 0) {
  25.                 result.add(node.val);
  26.                 continue;
  27.             }
  28.             if (node.left != null) {
  29.                 node.left.parent = null;
  30.                 q.add(node.left);
  31.                 qk.add(k - 1);
  32.             }
  33.             if (node.right != null) {
  34.                 node.right.parent = null;
  35.                 q.add(node.right);
  36.                 qk.add(k - 1);
  37.             }
  38.             if (node.parent != null) {
  39.                 if (node.parent.right != null && node.parent.right.val == node.val)
  40.                     node.parent.right = null;
  41.                 else
  42.                     node.parent.left = null;
  43.                 q.add(node.parent);
  44.                 qk.add(k - 1);
  45.             }
  46.         }
  47.  
  48.         return result;
  49.     }
  50.  
  51.     private ParentedTreeNode FindTargetNode(ParentedTreeNode treeNode, int val) {
  52.         Queue<ParentedTreeNode> q = new LinkedList<>();
  53.         q.add(treeNode);
  54.         while (!q.isEmpty()) {
  55.             ParentedTreeNode node = q.remove();
  56.             if (node.val == val)
  57.                 return node;
  58.             if (node.left != null)
  59.                 q.add(node.left);
  60.             if (node.right != null)
  61.                 q.add(node.right);
  62.         }
  63.         return null;
  64.     }
  65.  
  66.     private ParentedTreeNode constructParentedTree(TreeNode root) {
  67.         Queue<TreeNode> qt = new LinkedList<>();
  68.         Queue<ParentedTreeNode> qpt = new LinkedList<>();
  69.  
  70.         qt.add(root);
  71.         ParentedTreeNode rootqpt = new ParentedTreeNode(root.val);
  72.         qpt.add(rootqpt);
  73.         while (!qt.isEmpty()) {
  74.             TreeNode nodeT = qt.remove();
  75.             ParentedTreeNode nodePT = qpt.remove();
  76.             if (nodeT.left != null) {
  77.                 qt.add(nodeT.left);
  78.                 nodePT.left = new ParentedTreeNode(nodeT.left.val);
  79.                 nodePT.left.parent = nodePT;
  80.                 qpt.add(nodePT.left);
  81.             }
  82.             if (nodeT.right != null) {
  83.                 qt.add(nodeT.right);
  84.                 nodePT.right = new ParentedTreeNode(nodeT.right.val);
  85.                 nodePT.right.parent = nodePT;
  86.                 qpt.add(nodePT.right);
  87.             }
  88.         }
  89.         return rootqpt;
  90.     }
  91.  
  92.     public static TreeNode stringToTreeNode(String input) {
  93.         input = input.trim();
  94.         input = input.substring(1, input.length() - 1);
  95.         if (input.length() == 0) {
  96.             return null;
  97.         }
  98.  
  99.         String[] parts = input.split(",");
  100.         String item = parts[0];
  101.         TreeNode root = new TreeNode(Integer.parseInt(item));
  102.         Queue<TreeNode> nodeQueue = new LinkedList<>();
  103.         nodeQueue.add(root);
  104.  
  105.         int index = 1;
  106.         while (!nodeQueue.isEmpty()) {
  107.             TreeNode node = nodeQueue.remove();
  108.  
  109.             if (index == parts.length) {
  110.                 break;
  111.             }
  112.  
  113.             item = parts[index++];
  114.             item = item.trim();
  115.             if (!item.equals("null")) {
  116.                 int leftNumber = Integer.parseInt(item);
  117.                 node.left = new TreeNode(leftNumber);
  118.                 nodeQueue.add(node.left);
  119.             }
  120.  
  121.             if (index == parts.length) {
  122.                 break;
  123.             }
  124.  
  125.             item = parts[index++];
  126.             item = item.trim();
  127.             if (!item.equals("null")) {
  128.                 int rightNumber = Integer.parseInt(item);
  129.                 node.right = new TreeNode(rightNumber);
  130.                 nodeQueue.add(node.right);
  131.             }
  132.         }
  133.         return root;
  134.     }
  135.  
  136.     public static String integerArrayListToString(List<Integer> nums, int length) {
  137.         if (length == 0) {
  138.             return "[]";
  139.         }
  140.  
  141.         String result = "";
  142.         for (int index = 0; index < length; index++) {
  143.             Integer number = nums.get(index);
  144.             result += Integer.toString(number) + ", ";
  145.         }
  146.         return "[" + result.substring(0, result.length() - 2) + "]";
  147.     }
  148.  
  149.     public static String integerArrayListToString(List<Integer> nums) {
  150.         return integerArrayListToString(nums, nums.size());
  151.     }
  152.  
  153.     public static void main(String[] args) throws IOException {
  154.         BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
  155.         String line;
  156.         while ((line = in.readLine()) != null) {
  157.             TreeNode root = stringToTreeNode(line);
  158.             line = in.readLine();
  159.             TreeNode target = stringToTreeNode(line);
  160.             line = in.readLine();
  161.             int K = Integer.parseInt(line);
  162.  
  163.             List<Integer> ret = new AllNodesDistanceKinBinaryTree().distanceK(root, target, K);
  164.  
  165.             String out = integerArrayListToString(ret);
  166.  
  167.             System.out.print(out);
  168.         }
  169.     }
  170.  
  171.     public class ParentedTreeNode {
  172.         int val;
  173.         ParentedTreeNode left;
  174.         ParentedTreeNode right;
  175.         ParentedTreeNode parent;
  176.  
  177.         ParentedTreeNode(int x) {
  178.             val = x;
  179.         }
  180.     }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement