Leetcode: 839. Similar String Groups

Can we formulate the problem using graphs? Words are vertices and they are connected if similar. The problem becomes finding the number of connected components in the graph using Union-Find. Time and space complexity is \(O(n)\).

from typing import List


class Solution:
    def numSimilarGroups(self, strs: List[str]) -> int:
        N = len(strs)
        M = len(strs[0])
        p = [i for i in range(N)]

        def find(u):
            if p[u] == u:
                return u
            p[u] = find(p[u])
            return p[u]

        def union(u, v):
            gu = find(u)
            gv = find(v)
            p[gu] = gv

        def is_similar(i, j):
            cnt = 0
            for k in range(M):
                if cnt > 2:
                    return False
                if strs[i][k] != strs[j][k]:
                    cnt += 1
            return cnt == 0 or cnt == 2

        for i in range(N):
            for j in range(i + 1, N):
                if is_similar(i, j):
                    union(i, j)

        comps = set(find(i) for i in range(N))
        return len(comps)


assert Solution().numSimilarGroups(["tars", "rats", "arts", "star"]) == 2
assert Solution().numSimilarGroups(["omv", "ovm"]) == 1