Union Find implementation with Path Compression and Union by Size

union_find.py

#!/usr/bin/env python3

from typing import Any, Dict, Iterable, Tuple


class UnionFind:
	def __init__(self, vs: Iterable[Any] = (), es: Iterable[Tuple[Any, Any]] = ()) -> None:
		self.parent: Dict[Any, Any] = {}
		self.size: Dict[Any, int] = {}
		for v in vs:
			self.parent[v] = v
			self.size[v] = 1
		self.num_components = len(self.parent)
		for v, w in es:
			self.union(v, w)

	def make_set(self, v: Any) -> bool:
		if self.contains(v):
			return False
		else:
			self.parent[v] = v
			self.size[v] = 1
			self.num_components += 1
			return True

	def union(self, v: Any, w: Any) -> bool:
		v_root = self.find(v)
		w_root = self.find(w)

		if v_root == w_root:
			return False
		else:
			if self.size[v_root] < self.size[w_root]:
				self.parent[v_root] = w_root
				self.size[w_root] += self.size[v_root]
			else:
				self.parent[w_root] = v_root
				self.size[v_root] += self.size[w_root]
			self.num_components -= 1
			return True

	def find(self, v: Any) -> Any:
		if not self.contains(v):
			raise KeyError(v)

		v_ancestor = v
		while v_ancestor != self.parent[v_ancestor]:
			v_ancestor = self.parent[v_ancestor]
		v_root = v_ancestor

		v_ancestor = v
		while v_ancestor != v_root:
			v_ancestor_parent = self.parent[v_ancestor]
			self.parent[v_ancestor] = v_root
			v_ancestor = v_ancestor_parent

		return v_root

	def contains(self, v: Any) -> bool:
		return v in self.parent and v in self.size

	def connected(self, v: Any, w: Any) -> bool:
		v_root = self.find(v)
		w_root = self.find(w)
		return v_root == w_root

	def component_size(self, v: Any) -> int:
		v_root = self.find(v)
		return self.size[v_root]