apriori.py

# coding:utf-8
from random import randint, sample
from collections import Counter
from itertools import combinations


def gen_carts(items, n_cart):
    n_items = len(items)
    carts = []
    with open('carts.txt', 'w') as f:
        for i in xrange(n_cart):
            n_cart_item = randint(1, n_items)
            cart = sample(items, n_cart_item)
            carts.append(cart)
            f.write(','.join(cart) + '\n')
    return carts


def apriori(carts, min_sup, k, kset, turn):
    print '=' * 10
    if turn > k:
        return

    items = []
    for cart in carts:
        for _set in combinations(cart, turn):
            if turn == 1:
                items.append(frozenset(_set))
            elif turn > 1 and all([set(item) in kset for item in combinations(_set, turn - 1)]):
                items.append(frozenset(_set))

    kset = []
    for item, n in Counter(items).iteritems():
        if n > min_sup:
            print item, n
            kset.append(item)
    if kset:
        apriori(carts, min_sup, k, kset, turn + 1)


def mining():
    items = ('bread', 'milk', 'apple', 'tomato', 'toothbrush')
    carts = gen_carts(items, 10)
    min_sup = 3
    k = 3
    kset = None
    turn = 1
    apriori(carts, min_sup, k, kset, turn)


if __name__ == '__main__':
    mining()