EquivalentTrees task: add test cases, adjust solution to cover them

Author: andrei-punko

src/main/java/by/andd3dfx/tree/equivalent/EquivalentTrees.java

@@ -32,9 +32,11 @@ public List<Node> findEquivalentSubtrees(Node root) {
 
         // Fill `vocabulary` field of nodes
         fillNodeVocabulary(root);
+        // Fill `subtreeSize` field of nodes
+        fillSubTreeSize(root);
 
         // Build Set<Character> -> List<Node> map
-        Map<Set<Character>, List<Node>> voc2Nodes = new HashMap<>();
+        Map<Set<Character>, List<Node>> voc2NodesMap = new HashMap<>();
 
         Deque<Node> queue = new ArrayDeque<>();
         queue.add(root);
@@ -46,23 +48,23 @@ public List<Node> findEquivalentSubtrees(Node root) {
             if (current.right != null) {
                 queue.add(current.right);
             }
-            Set<Character> vocabulary = current.vocabulary;
 
-            if (!voc2Nodes.containsKey(vocabulary)) {
-                voc2Nodes.put(vocabulary, new ArrayList<>());
+            Set<Character> vocabulary = current.vocabulary;
+            if (!voc2NodesMap.containsKey(vocabulary)) {
+                voc2NodesMap.put(vocabulary, new ArrayList<>());
             }
-            voc2Nodes.get(vocabulary).add(current);
+            voc2NodesMap.get(vocabulary).add(current);
         }
 
-        return voc2Nodes.values().stream()
-                .filter(nodesList -> nodesList.size() >= 2)
-                .sorted((o1, o2) -> o2.stream().mapToInt(nodeToIntFunction()).sum() - o1.stream().mapToInt(nodeToIntFunction()).sum())
-                .limit(1)
-                .findFirst().orElse(null);
+        return voc2NodesMap.values().stream()
+                .filter(nodes -> nodes.size() >= 2)
+                .min((o1, o2) -> o2.stream().mapToInt(nodeToIntFunction()).sum() - o1.stream().mapToInt(nodeToIntFunction()).sum())
+                .map(nodes -> nodes.subList(0, 2))
+                .orElse(null);
     }
 
     private static ToIntFunction<Node> nodeToIntFunction() {
-        return node -> node.vocabulary.size();
+        return node -> node.subtreeSize;
     }
 
     private Set<Character> fillNodeVocabulary(Node node) {
@@ -76,4 +78,14 @@ private Set<Character> fillNodeVocabulary(Node node) {
         }
         return node.vocabulary;
     }
+
+    private int fillSubTreeSize(Node node) {
+        if (node.left != null) {
+            node.subtreeSize += 1 + fillSubTreeSize(node.left);
+        }
+        if (node.right != null) {
+            node.subtreeSize += 1 + fillSubTreeSize(node.right);
+        }
+        return node.subtreeSize;
+    }
 }

src/main/java/by/andd3dfx/tree/equivalent/Node.java

@@ -9,6 +9,7 @@ public class Node {
     Node right;
 
     Set<Character> vocabulary = new HashSet<>();
+    int subtreeSize;
 
     public Node(char value) {
         this.value = value;

src/test/java/by/andd3dfx/tree/equivalent/EquivalentTreesTest.java

@@ -79,7 +79,7 @@ public void findEquivalentSubtreesWhenOneCandidate() {
      * </pre>
      */
     @Test
-    public void findEquivalentSubtreesWhenTwoCandidate() {
+    public void findEquivalentSubtreesSymmetricCase() {
         Node node = new Node('A');
         node.left = new Node('B');
         node.left.left = new Node('X');
@@ -92,6 +92,58 @@ public void findEquivalentSubtreesWhenTwoCandidate() {
         assertThat(result, hasItems(node.left, node.right));
     }
 
+    /**
+     * <pre>
+     *      A
+     *     / \
+     *    B   B
+     *   /     \
+     *  X       C
+     *           \
+     *            X
+     * </pre>
+     */
+    @Test
+    public void findEquivalentSubtreesAsymmetricCase() {
+        Node node = new Node('A');
+        node.left = new Node('B');
+        node.left.left = new Node('X');
+        node.right = new Node('B');
+        node.right.right = new Node('C');
+        node.right.right.right = new Node('X');
+
+        List<Node> result = equivalentTrees.findEquivalentSubtrees(node);
+
+        assertThat(result.size(), is(2));
+        assertThat(result, hasItems(node.left, node.right.right));
+    }
+
+    /**
+     * <pre>
+     *      A
+     *     / \
+     *    B   B
+     *   /     \
+     *  X       X
+     *           \
+     *            X
+     * </pre>
+     */
+    @Test
+    public void findEquivalentSubtreesAsymmetricCaseShouldChooseNodesWithMaxSubtreeSize() {
+        Node node = new Node('A');
+        node.left = new Node('B');
+        node.left.left = new Node('X');
+        node.right = new Node('B');
+        node.right.right = new Node('X');
+        node.right.right.right = new Node('X');
+
+        List<Node> result = equivalentTrees.findEquivalentSubtrees(node);
+
+        assertThat(result.size(), is(2));
+        assertThat(result, hasItems(node.left, node.right));
+    }
+
     /**
      * <pre>
      *        A
@@ -101,6 +153,8 @@ public void findEquivalentSubtreesWhenTwoCandidate() {
      *    E   B   E
      *   /       / \
      *  D       E   D
+     *             / \
+     *            D   E
      * </pre>
      */
     @Test
@@ -116,6 +170,8 @@ public void findEquivalentSubtreesComplexCase() {
         root.right.right = new Node('E');
         root.right.right.left = new Node('E');
         root.right.right.right = new Node('D');
+        root.right.right.right.left = new Node('D');
+        root.right.right.right.right = new Node('E');
 
         List<Node> result = equivalentTrees.findEquivalentSubtrees(root);