Code I wrote in an interview with Google. Totally nailed it -- and then got ghosted for a month by the recruiter until he got back to me to say that the position was no longer open.

google-union-find.py

"""
You're given a list of elements. Each "element" has a unique id and 3 properties.

Two elements E1 and E2 are "similar" if they share any of their 3 properties.
Please write a function that takes a list of elements as input, and returns all "similar" groups of elements.

Example Input: [E1, E2, E3]
  E1: (id1, p1, p2, p3)
  E2: (id2, p1, p4, p5)
  E3: (id3, p6, p7, p8)

Correct output:
[ [E1, E2], [E3] ]

Note that:
  1. Similarity is transitive. If E1 ~ E2 and E2 ~ E3, then E1 ~ E3.
  2. The order of properties p does not matter
"""

from collections import defaultdict


class UnionFind:
    def __init__(self, enames):
        # key = element-name, value = parent
        self.groups = {ename: ename for ename in enames}
        # self.heads = set(self.groups) # implicitly keys of dict

    def parent_of(self, ename):
        while (parent := self.groups.get(ename)) != ename:
            ename = parent
        return parent

    def join(self, ename1, ename2):
        p1 = self.parent_of(ename1)
        p2 = self.parent_of(ename2)
        if p1 != p2:
            self.groups[p1] = p2
            # self.heads.pop(p1)


def find_similar(elements):
    # property to elements having this property
    d = defaultdict(set)
    for ename, eprops in elements:
        for eprop in eprops:
            d[eprop].add(ename)
    # d.values() has sets of related elements

    # Create union-find
    element_names = [ename for ename, _ in elements]
    uf = UnionFind(element_names)

    # Populate union-find with pairs of related elements
    for related_set in d.values():
        related_elements = list(related_set)
        for i in range(1, len(related_elements)):
            e1, e2 = related_elements[i - 1], related_elements[i]
            uf.join(e1, e2)
    #print(uf.groups)

    # Extract related groups from UF
    related_groups = defaultdict(set)
    for key in uf.groups:
        parent = uf.parent_of(key)
        related_groups[parent].add(key)
    #print(related_groups)
    result = [sorted(value) for value in related_groups.values()]
    #print(result)
    return result


if __name__ == '__main__':
    assert find_similar([]) == []
    assert find_similar([['E1', tuple()]]) == [["E1"]]
    assert find_similar([
        ['E1', ("id1", )],
        ['E2', ("id2", )],
    ]) == [["E1"], ["E2"]]
    assert find_similar([
        ['E1', ("id1", )],
        ['E2', ("id1", )],
    ]) == [["E1", "E2"]]
    assert find_similar([
        ['E1', ("id1", )],
        ['E2', ("id1", "id2", )],
        ['E3', ("id2", )],
    ]) == [["E1", "E2", "E3"]]
"""
  E1: (id2)
  E2: (id2, id3)
  E3: (id3)
  
  E1: (id1, id2)
  E2: (id2, id3)
  E3: (id3, id1)
 
  E1: (id1)
  E2: (id2)
  E3: (id3)

  []
  
  [E1]
    
  E1: (id1)
  E2: (id1)

  E1: (id1)
  E2: (id2)
"""