From 480d4cf8308ce51c62a7f20532d0cd56034eec77 Mon Sep 17 00:00:00 2001 From: Pragy Agarwal Date: Tue, 17 Mar 2020 17:13:40 +0530 Subject: [PATCH] add disjoint set --- DS_implementations/python/disjoint_set.py | 45 +++++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 DS_implementations/python/disjoint_set.py diff --git a/DS_implementations/python/disjoint_set.py b/DS_implementations/python/disjoint_set.py new file mode 100644 index 0000000..d2174c3 --- /dev/null +++ b/DS_implementations/python/disjoint_set.py @@ -0,0 +1,45 @@ +from collections import defaultdict +from typing import Generic, TypeVar, DefaultDict, Dict + +T = TypeVar('T') + + +class DisjointSet(Generic[T]): + """Generic Disjoint Set implementation""" + + def __init__(self): + self._parents: Dict[T, T] = {} + self._sizes: DefaultDict[T, int] = defaultdict(lambda: 1) + + def find_root(self, x: T) -> T: + """ + Find the component in which x belongs in. + Uses path compression. + Amortized Time complexity per query ~O(1) + """ + node, root = None, x + while root != node: # using loop instead of recursion due to Python recursion limit + node, root = root, self._parents.get(root, root) + while x != root: + x = self._parents.get(x, x) + self._parents[x] = root # path compression + return root + + def size_of(self, x: T) -> int: + """ + Find the size of the component that x belongs in + Time Complexity per query: O(1) + """ + return self._sizes[self.find_root(x)] + + def merge_components(self, x: T, y: T) -> None: + """ + Merge the components in which the items x and y belong in + Amortized Time Complexity per query: ~O(1) + """ + x, y = self.find_root(x), self.find_root(y) + if self.size_of(x) < self.size_of(y): + x, y = y, x # swap so that y is always the larger component + self._parents[y] = x + self._sizes[x] += self._sizes[y] + del self._sizes[y]