package com.jiucaiyuan.net.algrithm.set;



import java.util.HashMap;

import java.util.List;

import java.util.Stack;



/**

 * @Author jiucaiyuan  2022/3/16 23:15

 * @mail services@jiucaiyuan.net

 */



public class UnionFindDemo {

    public static class Element<V> {

        public V value;

        public Element(V value) {

            this.value = value;

        }

    }



    public static class UnionFindSet<V> {

        //每个对象对应的元素

        public HashMap<V, Element<V>> elementMap;

        // <某元素，该元素的父>   直接父元素，不一定是代表元素

        public HashMap<Element<V>, Element<V>> fatherMap;

        // <集合代表元素，集合大小>

        public HashMap<Element<V>, Integer> sizeMap;



        /**

         * 并查集构造函数，初始化集合中元素，每个元素为一个集合

         * 局限：并查集初始化时，要求指定集合中所有元素

         * @param list 指定初始化数据范围

         */

        public UnionFindSet(List<V> list) {

            elementMap = new HashMap<>();

            fatherMap = new HashMap<>();

            sizeMap = new HashMap<>();

            for (V value : list) {

                Element<V> element = new Element<>(value);

                elementMap.put(value, element);

                fatherMap.put(element, element);

                sizeMap.put(element, 1);

            }

        }



        //找到指定元素的代表元素

        //当调用的特别频繁时，可以认为时间复杂度是O(1)，如果调用次数少，时间复杂度到不了O(1)

        private Element<V> findHead(Element<V> element) {

            Stack<Element<V>> path = new Stack<>();

            //找到代表节点

            while (element != fatherMap.get(element)) {

                path.push(element);

                element = fatherMap.get(element);

            }

            //扁平化处理，把所有元素的代表节点记录下来

            while(!path.isEmpty()){

                fatherMap.put(path.pop(),element);

            }

            return element;

        }



        /**

         * 判断集合a和集合b是否是同一个集合

         * @param a

         * @param b

         * @return

         */

        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;

        }



        /**

         * 合并集合a和集合b

         * @param a

         * @param b

         */

        public void union(V a,V b){

            if(elementMap.containsKey(elementMap.get(a)) && elementMap.containsKey(elementMap.get(b))){

                Element<V> aHead = findHead(elementMap.get(a));

                Element<V> bHead = findHead(elementMap.get(b));

                //如果不是同一个集合，进行合并操作

                if(aHead != bHead){

                    Element<V> bigHead = sizeMap.get(aHead)>=sizeMap.get(bHead)?aHead:bHead;

                    Element<V> smallHead = bigHead == aHead?bHead:aHead;

                    //设置父节点即可，在findHead方法时，会把fatherMap的关系设置成代表元素

                    fatherMap.put(smallHead,bigHead);

                    sizeMap.put(bigHead,sizeMap.get(aHead) + sizeMap.get(bHead));

                    sizeMap.remove(smallHead);

                }

            }

        }

    }

}