随笔记录

算法-并查集

2022-3-16 diaba 数据结构

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);
                }
            }
        }
    }
}