0
点赞
收藏
分享

微信扫一扫

Java 并查集代码

A邱凌 2022-01-06 阅读 47
java
public class UnionFind<V> {

    private class Node<V> {
        V value;
        public Node(V v) {
            value = v;
        }
    }
    
    // 元素与节点对应关系
    private HashMap<V, Node<V>> nodes;
    // 节点与父节点对应关系
    private HashMap<Node<V>, Node<V>> parents;
    // 根节点与根节点所在集合大小的对应关系
    private HashMap<Node<V>, Integer> sizeMap;

    public UnionFind(List<V> values) {
        nodes = new HashMap<>();
        parents = new HashMap<>();
        sizeMap = new HashMap<>();
        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;
    }

    // a 与 b 是否在同一个集合
    public boolean isSameSet(V a, V b) {
        return findFather(nodes.get(a)) == findFather(nodes.get(b));
    }

    // 合并 a 与 b 所在的集合
    public void union(V a, V b) {
        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);
        }
    }

    // 并查集内集合个数
    public int size() {
        return sizeMap.size();
    }

}
举报

相关推荐

0 条评论