
package codesample.algorithms.union;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * This is a Union implementation with quick 'connect' operation.
 * Union (node1, node2). Takes uo to n
 * Connected (node1, node2). Takes up to n
 * This is not the implementation you should be using a lot. Should only
 * be used as a test subject when learning algorithms
 */
public class UnionWithQuickUnionOperation<T extends Comparable<T>> implements Union<T>{

    private Map<T, T> map;

    public UnionWithQuickUnionOperation() {
        map = new HashMap<>();
    }

    public UnionWithQuickUnionOperation(Set<T> set) {
        /* Entires refer to themselfs */
        this.map = set.stream()
                .collect(Collectors.toMap(value -> value, value -> value));
    }

    /**
     * Find the root of hierarchy. Root element has a mapping to itself in map
     */
    private T root(T value) {
        var currentValue = value;
        while (map.get(currentValue) != currentValue) {
            currentValue = map.get(currentValue);
        }
        return currentValue;
    }

    @Override
    public boolean areConnected(T nodeOne, T nodeTwo) {
        return root(map.get(nodeOne)).equals(
                root(map.get(nodeTwo)));
    }

    @Override
    public void connect(T nodeOne, T nodeTwo) {
        addUnlinkedNode(nodeOne);
        addUnlinkedNode(nodeTwo);

        T rootLeft = root(nodeOne);
        T rootRight = root(nodeTwo);
        map.put(rootLeft, rootRight);
    }

    @Override
    public void addUnlinkedNode(T node) {
        if (!map.containsKey(node))
            map.put(node, node);
    }

    @Override
    public void removeNode(T node) {

        if (!map.containsKey(node))
            return;

        /* Nodes which should be remapped to other nodes when current node is removed */
        var toBeRemapped = map.keySet().stream()
                .filter(key -> map.get(key).equals(node))
                .collect(Collectors.toList());

        /* If nothing to remap, simply remove the node */
        if (toBeRemapped.isEmpty())
            map.remove(node);

        /* If this node is a root node, all node beneath it should be remapped to a single common node (it can be any
        * node of the list but we simply pick the first one here). This node also become a root node.
        * e.g. for:
        *   A -> B
        *   C -> B
        *   D -> B
        *REMOVE (B)
        *  A -> A
        *  C -> A
        *  D -> A */
        if (map.get(node).equals(node)) {
            var singleCommonNode = toBeRemapped.get(0);
            map.put(singleCommonNode, singleCommonNode);
            toBeRemapped.forEach(value -> map.put(value, singleCommonNode));

        /* If node-to-be-removed is not root, all nodes beneath it should be remapped to the node which node-to-be-removed
        * is mapped to
        * eg for:
        *   A -> B -> C
        *   D -> B -> C
        * REMOVE(B)
        *   A -> C
        *   D -> C */
        } else {
            var upperNode = map.get(node);
            toBeRemapped.forEach(value -> map.put(value, upperNode));
        }

        map.remove(node);
    }
}
   

@@ -212,21 +212,23 @@
 ion.%0A * 
-Union
+connect
  (node1,
@@ -253,16 +253,19 @@
 to n%0A * 
+are
 Connecte
@@ -296,16 +296,19 @@
 up to n%0A
+ *%0A
  * This 
@@ -314,14 +314,14 @@
  is 
-not th
+a naiv
 e im
@@ -336,101 +336,52 @@
 tion
- you should be using a lot. Should only%0A * be used as a test subject when learning algorithms
+. A lot of improvement can be applied to it.
 %0A */
@@ -2430,16 +2430,17 @@
        *
+ 
 REMOVE (
@@ -2730,16 +2730,46 @@
 nNode));
+%0A            map.remove(node);
 %0A%0A      
@@ -3169,23 +3169,16 @@
 ));%0A    
-    %7D%0A%0A
         
@@ -3191,24 +3191,36 @@
 move(node);%0A
+        %7D%0A%0A%0A
     %7D%0A%7D%0A   %0A

Add better comments to UnionWithQuickUnion Operation