1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
| import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Stack;
public class UnionFind { public static class Element<V> { private V value; public Element(V v) { this.value = v; } }
public static class UnionFindSet<V> { public Map<V, Element<V>> elementMap; public Map<Element<V>, Element<V>> fatherMap; public Map<Element<V>, Integer> rankMap;
public UnionFindSet(List<V> list) { elementMap = new HashMap<>(); fatherMap = new HashMap<>(); rankMap = new HashMap<>();
for (V v : list) { Element<V> element = new Element<V>(v); elementMap.put(v, element); fatherMap.put(element, element); rankMap.put(element, 1); } }
public 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; }
public boolean isSameSet(V a, V b) { if (!elementMap.containsKey(a) || !elementMap.containsKey(b)) { return false; }
return findHead(elementMap.get(a)) == findHead(elementMap.get(b)); }
public void union(V a, V b) { if (elementMap.containsKey(a) && elementMap.containsKey(b)) { Element<V> aH = findHead(elementMap.get(a)); Element<V> bH = findHead(elementMap.get(b)); if (aH != bH) { Element<V> big = rankMap.get(aH) >= rankMap.get(bH) ? aH : bH; Element<V> small = big == aH ? bH : aH; fatherMap.put(small, big); rankMap.put(big, rankMap.get(big) + rankMap.get(small)); rankMap.remove(small); } } } } }
|