并查集的Java工具类实现

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiPredicate;
import java.util.function.Function;

/**
 * @author zhaozhen
 */
public class DisjoinSet<K, V> {

    /**
     * 用 Map 在存储并查集，表达的含义是 key 的父节点是 value
     */
    private final Map<K, K> disjoinMap;
    private final Map<K, V> kvMap;

    private final List<V> dataList;

    private final Function<V, K> buildKeyFunction;

    public DisjoinSet(List<V> dataList, Function<V, K> function) {
        this.dataList = dataList;
        buildKeyFunction = function;
        disjoinMap = new ConcurrentHashMap<>();
        kvMap = new HashMap<>();
        for (V v : dataList) {
            disjoinMap.put(function.apply(v), function.apply(v));
            kvMap.put(function.apply(v), v);
        }
    }

    /**
     * 查找：反复查找父亲节点。
     */
    synchronized public K findRoot(K x) {
        // 寻找x祖先节点保存到root中
        K root = x;

        while (!disjoinMap.get(root).equals(root)) {
            root = disjoinMap.get(root);
        }
        // 路径压缩，把x到root上所有节点都挂到root下面
        while (!x.equals(root)) {

            // 保存原来的父节点
            K originalFather = disjoinMap.get(x);
            // 当前节点挂到根节点下面
            disjoinMap.put(x, root);
            // x赋值为原来的父节点继续执行刚刚的操作
            x = originalFather;
        }

        return root;
    }

    /**
     * 合并：把两个集合合并为一个，只需要把其中一个集合的根节点挂到另一个集合的根节点下方
     */
    synchronized public void union(K x, K y) {
        K rootX = findRoot(x);
        K rootY = findRoot(y);

        if (!rootX.equals(rootY)) {
            // 节点联通只需要一个共同祖先，无所谓谁是根节点
            disjoinMap.put(rootX, rootY);
        }
    }

    /**
     * 判断：判断两个元素是否同属一个集合
     */
    public boolean isConnected(K x, K y) {
        return findRoot(x).equals(findRoot(y));
    }

    public Map<K, List<V>> getInsectGroup(BiPredicate<V, V> biPredicate) {
        for (int i = 0; i < dataList.size(); i++) {
            V dataI = dataList.get(i);
            for (int j = i + 1; j < dataList.size(); j++) {
                V dataJ = dataList.get(j);
                if (biPredicate.test(dataI, dataJ)) {
                    union(buildKeyFunction.apply(dataI), buildKeyFunction.apply(dataJ));
                }
            }
        }

        // 保证所有的树都只有两层
        for (int i = 0; i < dataList.size(); i++) {
            V dataI = dataList.get(i);
            findRoot(buildKeyFunction.apply(dataI));
        }


        Map<K, List<V>> resMap = new HashMap<>();
        for (Map.Entry<K, K> entry : disjoinMap.entrySet()) {
            K f = entry.getKey();
            K r = entry.getValue();
            if (!resMap.containsKey(r)) {
                List<V> list = new ArrayList<>();
                list.add(kvMap.get(f));
                resMap.put(r, list);
            } else {
                List<V> list = resMap.get(r);
                list.add(kvMap.get(f));
                resMap.put(r, list);
            }
        }
        return resMap;
    }

}