并查集

package newcoder;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class UnionFind {

    public static class Element<V> {
        public V value;
        
        public Element (V value) {
            this.value = value;
        }
    }
    
    public static class UnionFindSet<V> {
        /**
         * 根据用户的value映射为包装类
         */
        public HashMap<V, Element<V>> elementMap;
        /**
         * 每个element都有自己的parent，代表节点的parent为自己
         */
        public HashMap<Element<V>, Element<V>> parentMap;
        /**
         * 每个代表节点代表了多少节点（如果不是代表节点，则为0）
         */
        public HashMap<Element<V>, Integer> sizeMap;
        
        /**
         * 初始化操作
         */
        public UnionFindSet(List<V> list) {
            elementMap = new HashMap<>();
            parentMap = new HashMap<>();
            sizeMap = new HashMap<>();
            
            for (V v : list) {
                Element<V> ele = new Element<V>(v);
                elementMap.put(v, ele);
                parentMap.put(ele, ele);
                sizeMap.put(ele, 1);
            }
        }
        
        /**
         * 方法1：判断两个节点是否为同一个集合
         */
        public boolean isSameSet(V a, V b) {
            if(elementMap.containsKey(a) && elementMap.containsKey(b)) {
                return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
            }
            return false;
        }
        
        /**
         * 方法2：合并两个元素
         * @param a
         * @param b
         */
        public void union(V a, V b) {
            if(elementMap.containsKey(a) && elementMap.containsKey(b)) {
                Element<V> aHead = findHead(elementMap.get(a));
                Element<V> bHead = findHead(elementMap.get(b));
                if(aHead != bHead) {
                    Element<V> big = sizeMap.get(aHead) > sizeMap.get(bHead) ? aHead : bHead;
                    Element<V> small = big == aHead ? bHead : aHead;
                    parentMap.put(small, big);
                    sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
                    sizeMap.remove(small);
                }
            }
        }
        
        /**
         * 私有方法：获取代表节点并将节点扁平化
         * @param ele
         * @return
         */
        private Element<V> findHead(Element<V> ele) {
             Stack<Element<V>> stack = new Stack<>();
             while(ele != parentMap.get(ele)) {
                 stack.push(ele);
                 ele = parentMap.get(ele);
             }
             
             while(!stack.isEmpty()) {
                 parentMap.put(stack.pop(), ele);
             }
             return ele;     
        }
    }
    
}