Accounts merge with union-find on strings in Python

Why this matters

Many real programs have records that describe the same thing in different ways. A shop may have two customer records for one person, one with a work email and one with a home email. A website may store the same user twice after a login provider changed. Finding these links is a grouping problem. Two records belong together if they share a piece of information.

This lesson solves a classic version of that problem. You get a list of accounts. Each account has a name and one or more email addresses. Two accounts belong to the same person if they share at least one email address. Your job is to merge the accounts that belong together and return one row per person, with that person's emails in sorted order.

The method uses union-find, the structure from the earlier lesson 05-union-find-and-minimum-spanning-trees. The speed rules are in the earlier lesson on union by size and path halving. Here union-find works on the text of the emails, which Python calls strings. You need to know how to write a function and how to use lists, dicts and for loops. A dict is a table of key and value pairs. Each Python type used below is explained where it first appears.

The problem in plain words

Here is a small input. Each row is one account. The first item is the name, and the rest are emails. In Python each row is a list, and the whole input is a list of lists.

["John", "johnsmith@mail.com", "john_newyork@mail.com"]
["John", "johnsmith@mail.com", "john00@mail.com"]
["Mary", "mary@mail.com"]
["John", "johnnybravo@mail.com"]

Read the first two rows. Both contain johnsmith@mail.com, so they are the same person. Their emails together are johnsmith@mail.com, john_newyork@mail.com and john00@mail.com. Mary has one account with one email, so she stays alone. The last John has a different email and shares nothing with the first two rows. He is a different person, even though his name is also John.

The correct answer is one merged row for John Smith, one row for Mary, and one row for the other John. The program in this lesson prints this output:

[['John', 'john00@mail.com', 'john_newyork@mail.com', 'johnsmith@mail.com'], ['John', 'johnnybravo@mail.com'], ['Mary', 'mary@mail.com']]

Notice two things. The name is not a key. Two people named John stay separate, because the link is the email, not the name. Also, the emails inside each row are sorted, which is the output format the problem asks for.

The key idea: join emails, not accounts

The simplest way to find the groups is to treat each email as a node, which is a point that can be joined to other points. Union-find works on numbers, so every distinct email needs a number. Once you have numbers, the rule is short:

• Inside one account, every email belongs to the same person. So for each account, join all its emails into one group.
• Two accounts that share an email automatically end up in one group. The shared email is a node that both accounts joined, so both accounts' emails are now connected through it.

For the first two rows above, both rows contain the same johnsmith email, which is one node, so the two rows are linked already. The first row also joins john_newyork to johnsmith. The second row joins john00 to johnsmith. Now the three emails form one group, without any extra check between the accounts.

The steps are:

1. Give each distinct email a number, starting at 0.
2. Make a union-find structure with one element per email.
3. For each account, join each of its emails to the first email of that account.
4. For each email, find its root. Emails with the same root belong to the same person.
5. For each group, sort the emails, take the name of one of its accounts, and write one output row.

Step 1: give each email a number

A dict stores pairs of a key and a value, and it finds the value for a key quickly. Here the key is an email string and the value is a number. These lines create two empty dicts:

email_id = {}
email_name = {}

The braces in {} make an empty dict. Read it as "a table from email text to a whole number". Readers who know Java can think of a dict as a Map.

Each time the program meets a new email, it gives that email the next free number. The next free number is the count of entries so far, because numbers start at 0. The built-in function len returns that count.

for account in accounts:
    name = account[0]
    for email in account[1:]:
        if email not in email_id:
            email_id[email] = len(email_id)
            email_name[email] = name

account[0] is the name. account[1:] is a slice. A slice is a new list with the items from index 1 to the end, so it holds only the emails. The not in test checks the dict's keys, so a repeated email does not get a second number.

Python works out the right side of an assignment first. So len(email_id) counts the entries before the new email is added. The first new email gets 0, the second gets 1, and so on.

The same loop saves the name of the first account that mentions each email, in a second dict called email_name. You need it at the end to label each group.

Step 2: the union-find structure

Once every email has a number, the number of elements is len(email_id), which gives uf = UnionFind(len(email_id)). The structure is a class named UnionFind. A class is a blueprint for an object that keeps its own data. The method __init__ runs when you create the object. self is the object itself, so self.parent is the list stored inside it.

class UnionFind:
    def __init__(self, size):
        self.parent = list(range(size))
 
    def find(self, x):
        parent = self.parent
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
 
    def union(self, a, b):
        ra = self.find(a)
        rb = self.find(b)
        if ra != rb:
            self.parent[rb] = ra

range(size) gives the numbers from 0 up to size - 1. list turns them into a list. So parent[i] starts as i, and each email starts as its own group. The list is sized by the number of distinct emails, not by the number of accounts. This choice matters, and the common mistakes section shows what happens when it is wrong.

find walks up the parent links until it reaches a root, which is the item whose parent is itself. The line parent[x] = parent[parent[x]] makes each item point one step higher as the walk goes. This is called path halving. It keeps later finds short.

union finds the root of both emails. If the roots differ, it makes one root the parent of the other. If the roots are the same, the two emails are already in one group, and nothing changes.

This version has no size rule. For this problem the input is small, so the simpler code is enough. For millions of records, add union by size from the earlier lesson.

Step 3: join the emails of each account

For each account, take its first email as the anchor. Join every other email in that account to the anchor.

for account in accounts:
    first = email_id[account[1]]
    for email in account[2:]:
        uf.union(first, email_id[email])

email_id[account[1]] looks up the number of the first email. The key goes in square brackets. The problem guarantees that each account has at least one email, so account[1] exists.

Joining to one anchor is enough. If three emails are in one account, joining the second and third to the first puts all three in one group. A group is a set of elements that share a root, and the anchor's root becomes the root of the whole group.

Step 4: group the emails by root

Now each email has a root, and emails with the same root are one person. A dict from root number to a list of emails collects them.

from collections import defaultdict
 
groups = defaultdict(list)
for email, number in email_id.items():
    groups[uf.find(number)].append(email)

email_id.items() gives each key and value as a pair. defaultdict(list) is a dict that creates an empty list the first time you use a new key. A plain dict would raise KeyError for a new root.

Step 5: build the output rows

Each group becomes one row. The row starts with the name, then lists the emails in sorted order.

result = []
for emails in groups.values():
    emails.sort()
    result.append([email_name[emails[0]]] + emails)
result.sort(key=lambda row: row[1])

emails.sort() sorts the list in place, in alphabetical order for strings. The name comes from email_name, looked up for the first email of the sorted group. Every email in one group belongs to the same person, so the name of any of them is a reasonable label when the names agree.

[name] + emails joins two lists into one new list.

The problem allows the output rows in any order. The sort call with key sorts the rows by their second item, which is the first email, so the output is the same on every run. lambda row: row[1] is a short unnamed function. It takes a row and returns the item at index 1.

sort changes the list it is called on and returns None. That is why result.sort(...) stands on its own line.

The full program

This is the whole program. It is the version that produced the outputs in this lesson. Save it as accounts_merge.py and run it with python3 accounts_merge.py. The outputs were produced with Python 3.14.

The output is:

[['John', 'john00@mail.com', 'john_newyork@mail.com', 'johnsmith@mail.com'], ['John', 'johnnybravo@mail.com'], ['Mary', 'mary@mail.com']]
[['Ann', 'a@x.com', 'b@x.com', 'c@x.com', 'd@x.com'], ['John', 'e@x.com']]
[['David', 'David0@m.co', 'David1@m.co', 'David2@m.co'], ['David', 'David3@m.co', 'David4@m.co']]

Example 1, traced by hand

Work through the first example on paper before you trust the code. The numbers are given in the order the emails first appear.

The unions come from each account. The first account joins number 1 to number 0. The second account joins number 2 to number 0. Mary's account and the last John's account have one email each, so they join nothing.

After the unions, the groups are {0, 1, 2}, {3} and {4}. The group {0, 1, 2} has three emails, so it becomes the merged John Smith row. The group {3} is Mary. The group {4} is the other John. Sorting the first group gives john00@mail.com, john_newyork@mail.com, johnsmith@mail.com, which matches the output.

Example 2: a chain across three accounts

The second example is harder because two accounts do not share an email directly. Here is the input:

["Ann", "a@x.com", "b@x.com"]
["Ann", "b@x.com", "c@x.com"]
["Ann", "c@x.com", "d@x.com"]
["John", "e@x.com"]

The first account links a and b. The second links b and c. The third links c and d. No single account contains both a and d, yet the chain joins all four. The union for each account is one step of the chain, and find follows the steps. The numbers are assigned in order: a is 0, b is 1, c is 2, d is 3, and e is 4.

Each account's union joins the new email to the anchor. The first account joins b to a. The second joins c to b. The third joins d to c. After these three joins, all four emails share one root, so they form one group. The last account, John with e@x.com, is alone, so it stays as its own row.

The output row for Ann lists a@x.com, b@x.com, c@x.com, d@x.com. This case shows why the method works with emails, not accounts. A single pass over the accounts would miss the link between a and d. Union-find keeps the link, because each join is remembered.

Example 3: duplicates and two people with one name

The third example has two people named David and an email that appears in several accounts. The first and last accounts share David1@m.co. The second and third accounts share David3@m.co and David4@m.co.

["David", "David0@m.co", "David1@m.co"]
["David", "David3@m.co", "David4@m.co"]
["David", "David4@m.co", "David3@m.co"]
["David", "David1@m.co", "David2@m.co"]

The emails David3@m.co and David4@m.co appear twice, once in each order. The second account joins them, and the third account joins them again. The second join finds that they are already in one group and does nothing. The code handles this safely, because union changes nothing when the two roots are equal.

The first and fourth accounts join David0, David1 and David2. The result has two rows, both named David, because the two groups share no email. Had the program merged by name, all four accounts would become one person with five emails. Matching on names is the wrong key, and the output above shows why.

What the method costs

Building the dict reads every email once, so that step takes time in proportion to the total number of email entries. Path halving keeps each find short in practice. Sorting each group takes a little more than linear time in the group size, but the groups are small in practice. The main cost is usually reading the input.

Memory is the other cost. The method keeps the dicts, the parent list and the groups, so it uses space in proportion to the number of distinct emails.

Common mistakes

List sized by the wrong count. The parent list must have one slot per distinct email. If you size it by the number of accounts, find reads an index past the end. Creating the structure with UnionFind(len(example1) - 1) on the first example (4 accounts, 5 emails) stops with this line at the end of the traceback:

IndexError: list index out of range

The traceback above that line points to while parent[x] != x: inside find, because find reads an index that the list does not have. The fix is to size the structure with len(email_id), which is the number of distinct emails.

Using an email as an index. A list index must be a whole number. An email is text, so this line fails:

parent = [0, 1]
parent["a@x.com"]

TypeError: list indices must be integers or slices, not str

The fix is to look up the number first, with email_id[email], and use that number as the index.

Looking up an email that is not in the dict. A dict raises KeyError for a missing key. It does not return None the way some languages do.

email_id = {"a@x.com": 0}
email_id["b@x.com"]

KeyError: 'b@x.com'

If you use get instead, a missing key returns None. Using that None in arithmetic then fails:

email_id.get("b@x.com") + 1

TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

The fix is to make sure every email from the accounts was added to the dict in step 1. The first loop must cover every account before any union runs.

Comparing strings with is. The == operator compares the text of two strings. The is operator checks whether two names point to the same object in memory. Two strings with the same text can be different objects:

a = "a@x.com"
b = "".join(["a@", "x.com"])
print(a == b, a is b)

True False

The == check is True, because the text is the same. The is check is False, because b was built at run time as a new object. Python sometimes shares one object for equal strings, so the result of is can change from one case to another. Use == to compare text.

Merging by name. A name is not a unique key. Two people can share a name, so joining accounts by name joins two different people. Example 3 shows the result of the right method. Keep emails as the only join key.

Practice

Exercise 1. Run accounts_merge on these accounts: ["Zoe", "z1@m.com"], ["Zoe", "z2@m.com"], and ["Zoe", "z1@m.com", "z2@m.com"]. How many rows are in the output, and what is in them?

Show the answer

There is one row: ['Zoe', 'z1@m.com', 'z2@m.com']. The first account has z1, the second has z2, and the third has both. The third account joins z2 to z1, which puts the two single-email accounts into one group. The program prints [['Zoe', 'z1@m.com', 'z2@m.com']].

Exercise 2. Run the method on ["A", "x@m.com"], ["B", "y@m.com"], and ["C", "x@m.com", "y@m.com"]. Which name appears in the output row, and why?

Show the answer

The output is [['A', 'x@m.com', 'y@m.com']]. The name is A, because x@m.com first appears in the account for A, and email_name keeps the name of the first account that mentions each email. The third account has the name C, but it does not decide the label. The problem does not say which name to use when accounts disagree, so state the rule you chose.

Exercise 3. An account lists the same email twice: ["Lee", "l@x.com", "l@x.com"]. Predict the output, and say what the union step does for this account.

Show the answer

The output is [['Lee', 'l@x.com']]. The email appears once in the output, because the dict gives it one number. The union step joins the number to itself. find returns the same root for both, so the if ra != rb check is false and nothing changes.

Exercise 4. The line if email not in email_id: in step 1 is removed, so every email is added again with email_id[email] = len(email_id). Explain what goes wrong, and run the code to check.

Show the answer

Assigning to a key that already exists replaces its value. So the repeated email gets the current length of the dict as its new number. In a test with the accounts ["Ann", "a@x.com", "b@x.com"] and ["Ann", "a@x.com", "a@x.com"], the second account gives a@x.com the number 2, which is past the end of the two-slot parent list. The program stops with IndexError: list index out of range. The check keeps one number for each email, so removing it breaks the core idea of the method.

Exercise 5. Change the program so it prints the number of people, not the rows. Write the one-line change and say what it prints for the first example.

Show the answer

Print len(accounts_merge(example1)) instead of accounts_merge(example1). The returned list has one row per person. For the first example the output is 3, because there are three rows: John Smith, Mary, and the other John.

Summary

• Give each distinct email a number with a dict, then build a union-find structure with one slot per email.
• For each account, join all its emails to the first email of that account.
• Group the emails by their root, and write one row per group with the emails sorted.
• Use emails as the only join key. Names can repeat between different people.
• Size the parent list by the number of distinct emails, and compare text with ==.