Leetcode 721 Solution

class FindUnionSet:
    def __init__(self, n):
        self.n = n
        self.p = list(range(n))

    def find(self, i):
        if self.p[i] == i:
            return i

        self.p[i] = self.find(self.p[i])
        return self.p[i]

    def union(self, i, j):
        ri = self.find(i)
        rj = self.find(j)

        self.p[ri] = rj


class Solution:
    def accountsMerge(self, accounts: List[List[str]]) -> List[List[str]]:
        mail_map = collections.defaultdict(list)

        for i, account in enumerate(accounts):
            for mail in range(1, len(account)):
                mail_map[account[mail]].append(i)

        fmset = FindUnionSet(len(accounts))
        for account_ids in mail_map.values():
            for i in range(1, len(account_ids)):
                fmset.union(account_ids[i - 1], account_ids[i])

        groups = collections.defaultdict(list)
        for i in range(len(accounts)):
            groups[fmset.find(i)].append(i)

        ans = []
        for account_ids in groups.values():
            account_name = None
            account_mails = set()
            for account_id in account_ids:
                account_name = accounts[account_id][0]
                for i in range(1, len(accounts[account_id])):
                    account_mails.add(accounts[account_id][i])

            ans.append([account_name] + sorted(list(account_mails)))

        return ans