LeetCode 721: Accounts Merge

UnionFind

We can build an email graph where vertices are the unique emails and an edge between two vertices exists if the emails belong to the same person. We can use UnionFind to keep track of the connected components (cc) of this email graph. In the end, we emit the (sorted) cc’s along with the owner name. If two emails appear in the same account, they belong to the same person, so, we union() them. If an email is common in two different accounts, the uniqueness of email-vertices ensures their connected components get merged.

With $n$ unique emails, time: $\mathcal{O}(n \cdot \lg^*{n} + n \cdot \lg{n})$ , space: $\mathcal{O}(n)$ .

class UnionFind:
    def __init__(self):
        self.parent_of, self.rank_of = {}, {}

    def __compress_path(self, u, root):
        while u != (p := self.parent_of[u]):
            self.parent_of[u] = root
            u = p

    def __find(self, u):
        root = u
        while root != (p := self.parent_of[root]):
            root = p
        self.__compress_path(u, root)
        return root

    def union(self, u, v):
        root_u = self.__find(u)
        root_v = self.__find(v)
        if root_u == root_v:
            return
        
        rank_u = self.rank_of[root_u]
        rank_v = self.rank_of[root_v]
        if rank_u == rank_v:
            self.parent_of[root_u] = root_v
            self.rank_of[root_v] += 1
            return
        
        if rank_u < rank_v:
            self.parent_of[root_u] = root_v
        else:
            self.parent_of[root_v] = root_u

    def add(self, u):
        self.parent_of[u] = u
        self.rank_of[u] = 0
    
    def connected_components(self):
        comp_map = {}
        for u in self.parent_of:
            root_u = self.__find(u)
            if root_u not in comp_map:
                comp_map[root_u] = []
            comp_map[root_u].append(u)
        
        return comp_map.values()


class Solution:
    def accountsMerge(self, accounts: List[List[str]]) -> List[List[str]]:
        email_name_map = {}

        uf = UnionFind()
        for acc in accounts:
            name = acc[0]
            for i in range(1, len(acc)):
                curr_email = acc[i]

                if curr_email not in email_name_map:
                    email_name_map[curr_email] = name
                    uf.add(curr_email)

                if i > 1:
                    prev_email = acc[i-1]
                    uf.union(prev_email, curr_email)
        
        merged_accounts = []
        for cc in uf.connected_components():
            account = []
            account.append( email_name_map[cc[0]] )
            account.extend(sorted(cc))
            merged_accounts.append(account)

        return merged_accounts

unreasonably effective