代码
public class UnionFind {
public static class Node<V>{
V value;
public Node(V v) {
value = v;
}
}
public static class UnionSet<V>{
public HashMap<V, Node<V>> nodes;
public HashMap<Node<V>, Node<V>> parents;
public HashMap<Node<V>, Integer> sizeMap;
public UnionSet(List<V> values) {
nodes = new HashMap<V,Node<V>>();
parents = new HashMap<Node<V>,Node<V>>();
sizeMap = new HashMap<Node<V>,Integer>();
for(V cur: values) {
Node<V> node = new Node<>(cur);
nodes.put(cur, node);
parents.put(node,node);
sizeMap.put(node, 1);
}
}
public Node<V> findFather(Node<V> cur){
Stack<Node<V>> path = new Stack<>();
while(cur != parents.get(cur)) {
path.push(cur);
cur = parents.get(cur);
}
while(!path.isEmpty()) {
parents.put(path.pop(),cur);
}
return cur;
}
public boolean isSameSet(V a,V b) {
if(!nodes.containsKey(a) || !nodes.containsKey(b)) {
return false;
}
return findFather(nodes.get(a)) == findFather(nodes.get(b));
}
public void union(V a,V b) {
if(!nodes.containsKey(a) || !nodes.containsKey(b)) {
return;
}
Node<V> aHead = findFather(nodes.get(a));
Node<V> bHead = findFather(nodes.get(b));
if(aHead != bHead) {
int aSetSize = sizeMap.get(aHead);
int bSetSize = sizeMap.get(bHead);
Node<V> big = aSetSize >= bSetSize ? aHead:bHead;
Node<V> small = big == aHead ? bHead:aHead;
parents.put(small,big);
sizeMap.put(big, aSetSize + bSetSize);
sizeMap.remove(small);
}
}
}
}