accountsMerge

accountsMerge.py

from collections import defaultdict

class Solution:
    
    def accountsMerge(self, accounts):
        """
        :type accounts: List[List[str]]
        :rtype: List[List[str]]
        """
        graph = defaultdict(set)
        mail_to_name = {}
        for acc in accounts:
            mail = acc[1]
            for m in acc[1:]:
                graph[mail].add(m)
                graph[m].add(mail)
                mail_to_name[m] = acc[0]
            
        result = []
        keys = set(graph.keys())
        for k in keys:
            if k not in graph:
                continue
            mails_k = graph.pop(k)
            stack = list(mails_k)
            while stack:
                for mail_j in graph.pop(stack.pop(), []):
                    if mail_j not in mails_k:
                        stack.append(mail_j)
                        mails_k.add(mail_j)
            
            result.append([mail_to_name[k]] + sorted(mails_k))
                
        return result