Disjoint Set Union in Python3

dsu.erl

-module(dsu).

-export([
	new/0,
	find/2,
	insert/2,
	union/3,
	groups/1
]).

-record(dsu, {
	parent = #{} :: #{term() => term()},
	rank = #{} :: #{term() => non_neg_integer()},
	count = 0 :: non_neg_integer()
}).

new() ->
	#dsu{}.

find(DSU = #dsu{parent = Parent}, X) when is_map_key(X, Parent) andalso map_get(X, Parent) =:= X ->
	{ok, X, DSU};
find(DSU0 = #dsu{parent = Parent0}, X) when is_map_key(X, Parent0) ->
	ParentX0 = map_get(X, Parent0),
	{ok, ParentX1, DSU1 = #dsu{parent = Parent1}} = find(DSU0, ParentX0),
	DSU2 = DSU1#dsu{parent = Parent1#{X => ParentX1}},
	{ok, ParentX1, DSU2}.

insert(DSU = #dsu{parent = Parent}, X) when is_map_key(X, Parent) ->
	{ok, DSU};
insert(DSU0 = #dsu{parent = Parent0, rank = Rank0, count = Count0}, X) ->
	DSU1 = DSU0#dsu{
		parent = Parent0#{X => X},
		rank = Rank0#{X => 0},
		count = Count0 + 1
	},
	{ok, DSU1}.

union(DSU = #dsu{parent = Parent}, X, Y) when is_map_key(X, Parent) andalso is_map_key(Y, Parent) andalso map_get(X, Parent) =:= map_get(Y, Parent) ->
	{ok, DSU};
union(DSU0 = #dsu{parent = Parent0}, X, Y) when is_map_key(X, Parent0) andalso is_map_key(Y, Parent0) ->
	{ok, RootX, DSU1} = find(DSU0, X),
	{ok, RootY, DSU2 = #dsu{parent = Parent1, rank = Rank0, count = Count0}} = find(DSU1, Y),
	case RootX =:= RootY of
		true ->
			{ok, DSU2};
		false when map_get(RootX, Rank0) =:= map_get(RootY, Rank0) ->
			Parent2 = Parent1#{RootX => RootY},
			Rank1 = Rank0#{RootX => map_get(RootX, Rank0) + 1},
			Count1 = Count0 - 1,
			DSU3 = DSU2#dsu{parent = Parent2, rank = Rank1, count = Count1},
			{ok, DSU3};
		false when map_get(RootX, Rank0) < map_get(RootY, Rank0) ->
			Parent2 = Parent1#{RootY => RootX},
			Count1 = Count0 - 1,
			DSU3 = DSU2#dsu{parent = Parent2, count = Count1},
			{ok, DSU3};
		false ->
			Parent2 = Parent1#{RootX => RootY},
			Count1 = Count0 - 1,
			DSU3 = DSU2#dsu{parent = Parent2, count = Count1},
			{ok, DSU3}
	end.

groups(DSU0 = #dsu{parent = Parent}) ->
	do_groups(DSU0, maps:iterator(Parent), #{}).

%% @private
do_groups(DSU0, Iterator0, Acc) ->
	case maps:next(Iterator0) of
		{K, _, Iterator1} ->
			{ok, ParentK, DSU1} = find(DSU0, K),
			do_groups(DSU1, Iterator1, Acc#{ParentK => []});
		none ->
			{ok, Acc, DSU0}
	end.

dsu.py

# Disjoint Set Union
# O(M * N) space
class UnionFind:
    parent: Dict[int, int]
    rank: Dict[int, int]
    count: int
    
    def __init__(self):
        self.parent = dict()
        self.rank = dict()
        self.count = 0
    
    # O(α(M * N)) time | O(1) space
    def find(self, x: int) -> int:
        while x != self.parent[x]:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x
    
    # O(1) time | O(1) space
    def insert(self, x: int) -> None:
        if x not in self.parent:
            self.parent[x] = x
            self.rank[x] = 0
            self.count += 1
    
    # O(α(M * N)) time | O(1) space
    def union(self, x: int, y: int) -> None:
        x = self.find(x)
        y = self.find(y)
        if x != y:
            if self.rank[x] < self.rank[y]:
                x, y = y, x
            self.parent[x] = y
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1
            self.count -= 1