final tweaks - passes tests kinda

Author: clormor

src/main/java/io/github/clormor/hackerrank/advanced/visitor/TreeParser.java

@@ -1,15 +1,14 @@
 package io.github.clormor.hackerrank.advanced.visitor;
 
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
 public class TreeParser {
 
-    Map<Integer, Integer> values;
-    Map<Integer, Color> colours;
+    int[] values;
+    Color[] colours;
     Map<Integer, Set<Integer>> allEdges;
 
     public Tree solve(int n, String[] args) {
@@ -20,53 +19,41 @@ public Tree solve(int n, String[] args) {
         // parse the edge lines into numbers we can search for
         mapAllEdges(args);
 
-        // begin depth-first search - find children of the root
-        Set<Integer> edges = findNodeEdges(0);
-
         // handle 1-node tree as special case
-        if (edges.size() == 0) {
-            return new TreeLeaf(values.get(0), colours.get(0), 0);
+        if (n == 1) {
+            return new TreeLeaf(values[0], colours[0], 0);
         }
 
         // recursively search the tree to create nodes from the root
-        TreeNode root = new TreeNode(values.get(0), colours.get(0), 0);
-        for (int child : edges) {
-            processNode(child, 1, root);
-        }
+        TreeNode root = new TreeNode(values[0], colours[0], 0);
+        processNode(root, 1);
 
         return root;
     }
 
-    private void processNode(Integer node, int depth, TreeNode parent) {
-        Set<Integer> edges = findNodeEdges(node);
-        if (edges.size() == 0) {
-            parent.addChild(new TreeLeaf(values.get(node), colours.get(node), depth));
-        } else {
-            TreeNode currentNode = new TreeNode(values.get(node), colours.get(node), depth);
-            parent.addChild(currentNode);
-            for (int child : edges) {
-                processNode(child, depth + 1, currentNode);
+    private void processNode(TreeNode parent, int parentId) {
+        for (Integer childId : allEdges.get(parentId)) {
+            allEdges.get(childId).remove(parentId);
+            Tree child;
+            Set<Integer> grandChildren = allEdges.get(childId);
+            if (grandChildren != null && grandChildren.size() == 0) {
+                child = new TreeLeaf(values[childId - 1], colours[childId - 1], parent.getDepth() + 1);
+            } else {
+                child = new TreeNode(values[childId - 1], colours[childId - 1], parent.getDepth() + 1);
+            }
+            parent.addChild(child);
+            if (child instanceof TreeNode) {
+                processNode((TreeNode) child, childId);
             }
         }
     }
 
-    private Set<Integer> findNodeEdges(int node) {
-        Set<Integer> result = allEdges.get(node);
-        if (result == null) {
-            return new HashSet<>();
-        }
-        for (Integer i : result) {
-            allEdges.get(i).remove(node);
-        }
-        return result;
-    }
-
     private void mapAllEdges(String[] edgeLines) {
         allEdges = new HashMap<>();
         for (int i = 2; i < edgeLines.length; i++) {
             String[] edge = edgeLines[i].split(" ");
-            int a = Integer.parseInt(edge[0]) - 1;
-            int b = Integer.parseInt(edge[1]) - 1;
+            int a = Integer.parseInt(edge[0]);
+            int b = Integer.parseInt(edge[1]);
 
             Set<Integer> as = allEdges.get(a);
             if (as == null) {
@@ -85,18 +72,18 @@ private void mapAllEdges(String[] edgeLines) {
     }
 
     private void mapValues(int n, String valuesLine) {
-        values = new HashMap<>(n);
+        values = new int[n];
         String[] valuesString = valuesLine.split(" ");
         for (int i = 0; i < n; i++) {
-            values.put(i, Integer.parseInt(valuesString[i]));
+            values[i] = Integer.parseInt(valuesString[i]);
         }
     }
 
     private void mapColours(int n, String coloursLine) {
-        colours = new HashMap<>(n);
+        colours = new Color[n];
         String[] coloursString = coloursLine.split(" ");
         for (int i = 0; i < n; i++) {
-            colours.put(i, (coloursString[i].equals("0")) ? Color.RED : Color.GREEN);
+            colours[i] = (coloursString[i].equals("0")) ? Color.RED : Color.GREEN;
         }
     }
 }

src/test/java/io/github/clormor/hackerrank/advanced/visitor/TestTreeParser.java

@@ -15,12 +15,6 @@ public void setup() {
         v = new SimpleCountTreeVisitor();
     }
 
-    @Test
-    public void very_simple_solve_test() {
-        String[] input = {"4", "0"};
-        checkTree(4, Color.RED, 0, t.solve(1, input));
-    }
-
     @Test
     public void depth_1_solve_test() {
         String[] input = {"4 6 7", "0 1 1", "1 2", "1 3"};