collab.py

import collections
import numpy as np

# customer -> product mapping
purchases = {
    'A': (0, 1),
    'B': (0, 2, 3),
    'C': (0, 3),
    'D': (1,),
    'E': (0, 1, 3),
    'F': (0, 3),
}

inverted_index = collections.defaultdict(list)
for customer, products in purchases.items():
    for p in products:
        inverted_index[p].append(customer)

# 4 products:
together = np.zeros((4, 4))

for product, customers in inverted_index.items():
    for c in customers:
        for product2 in purchases[c]:
            if product != product2:
                together[product, product2] += 1

for p in xrange(4):
    print "people who bought {} also bought {}".format(p, np.argsort(together[p,:])[::-1][:2])