Last Updated: February 25, 2016
·
459
· grunskis

How I learned to love python itertools

Imagine that you have a list of user IDs:

user_ids = [1, 2, 3]

and you need to check if these users are friends with each other and if not then make them friends.

There are 9 different combinations to check:

combinations = [
    (1, 1)
    (1, 2)
    (1, 3)
    (2, 1)
    (2, 2)
    (2, 3)
    (3, 1)
    (3, 2)
    (3, 3)
]

We can exclude most of the combinations because user cannot be friends with oneself and we can assume that if user A is friends with user B then also user B is friends with user A. So we're left with these combinations to check:

combinations = [
    (1, 2)
    (1, 3)
    (2, 3)
]

Simplest solution that comes into mind is write nested for loops:

for a in user_ids:
    for b in user_ids:
        if a == b:
            # user cannot be friends with oneself
            continue

        if a > b:
            # exclude duplicate/reverse combinations
            continue

        if not friends(a, b):
            create_friendship(a, b)

Let's try to simplify this code by using python itertools library:

import itertools

for a, b in itertools.combinations(user_ids, 2):
    if not friends(a, b):
        create_friendship(a, b)

\o/