modify code

master
algorithmzuo 2 years ago
parent af099345f2
commit c0e9fe57f2

@ -14,64 +14,59 @@ import java.util.Stack;
// 一定要看! // 一定要看!
public class Code05_UnionFind { public class Code05_UnionFind {
public static class Node<V> { // 课上讲的时候
V value; // 包了一层
// 其实不用包一层哦
public Node(V v) {
value = v;
}
}
public static class UnionFind<V> { public static class UnionFind<V> {
public HashMap<V, Node<V>> nodes; public HashMap<V, V> father;
public HashMap<Node<V>, Node<V>> parents; public HashMap<V, Integer> size;
public HashMap<Node<V>, Integer> sizeMap;
public UnionFind(List<V> values) { public UnionFind(List<V> values) {
nodes = new HashMap<>(); father = new HashMap<>();
parents = new HashMap<>(); size = new HashMap<>();
sizeMap = new HashMap<>();
for (V cur : values) { for (V cur : values) {
Node<V> node = new Node<>(cur); father.put(cur, cur);
nodes.put(cur, node); size.put(cur, 1);
parents.put(node, node);
sizeMap.put(node, 1);
} }
} }
// 给你一个节点,请你往上到不能再往上,把代表返回 // 给你一个节点,请你往上到不能再往上,把代表返回
public Node<V> findFather(Node<V> cur) { public V findFather(V cur) {
Stack<Node<V>> path = new Stack<>(); Stack<V> path = new Stack<>();
while (cur != parents.get(cur)) { while (cur != father.get(cur)) {
path.push(cur); path.push(cur);
cur = parents.get(cur); cur = father.get(cur);
} }
while (!path.isEmpty()) { while (!path.isEmpty()) {
parents.put(path.pop(), cur); father.put(path.pop(), cur);
} }
return cur; return cur;
} }
public boolean isSameSet(V a, V b) { public boolean isSameSet(V a, V b) {
return findFather(nodes.get(a)) == findFather(nodes.get(b)); return findFather(a) == findFather(b);
} }
public void union(V a, V b) { public void union(V a, V b) {
Node<V> aHead = findFather(nodes.get(a)); V aFather = findFather(a);
Node<V> bHead = findFather(nodes.get(b)); V bFather = findFather(b);
if (aHead != bHead) { if (aFather != bFather) {
int aSetSize = sizeMap.get(aHead); int aSize = size.get(aFather);
int bSetSize = sizeMap.get(bHead); int bSize = size.get(bFather);
Node<V> big = aSetSize >= bSetSize ? aHead : bHead; if (aSize >= bSize) {
Node<V> small = big == aHead ? bHead : aHead; father.put(bFather, aFather);
parents.put(small, big); size.put(aFather, aSize + bSize);
sizeMap.put(big, aSetSize + bSetSize); size.remove(bFather);
sizeMap.remove(small); } else {
father.put(aFather, bFather);
size.put(bFather, aSize + bSize);
size.remove(aFather);
}
} }
} }
public int sets() { public int sets() {
return sizeMap.size(); return size.size();
} }
} }

Loading…
Cancel
Save