Efficient Implementation of the Apriori Algorithm with Numpy

#! -*- coding: utf-8 -*-

import numpy as np


class Apriori:

    def __init__(self, min_support, min_confidence):
        self.min_support = min_support # Minimum support
        self.min_confidence = min_confidence # Minimum confidence

    def count(self, filename='apriori.txt'):
        self.total = 0 # Total number of data rows
        items = {} # Item list

        # Statistics to get the item list
        with open(filename) as f:
            for l in f:
                self.total += 1
                for i in l.strip().split(','): # Separated by commas
                    if i in items:
                        items[i] += 1.
                    else:
                        items[i] = 1.

        # Deduplicate item list and map to IDs
        self.items = {i:j/self.total for i,j in items.items() if j/self.total > self.min_support}
        self.item2id = {j:i for i,j in enumerate(self.items)}

        # 0-1 matrix of the item list
        self.D = np.zeros((self.total, len(items)), dtype=bool)

        # Re-traverse the file to get the 0-1 matrix
        with open(filename) as f:
            for n,l in enumerate(f):
                for i in l.strip().split(','):
                    if i in self.items:
                        self.D[n, self.item2id[i]] = True

    def find_rules(self, filename='apriori.txt'):
        self.count(filename)
        rules = [{(i,):j for i,j in self.items.items()}] # Record frequent itemsets at each step
        l = 0 # Number of items in the frequent itemsets of the current step

        while rules[-1]: # Includes the construction process from k-frequent to k+1 frequent itemsets
            rules.append({})
            keys = sorted(rules[-2].keys()) # Sort each k-frequent itemset lexicographically (Core)
            num = len(rules[-2])
            l += 1
            for i in range(num): # Traverse each pair of k-frequent itemsets
                for j in range(i+1,num):
                    # If the first k-1 items overlap, these two k-frequent itemsets can be combined into a k+1 frequent itemset
                    if keys[i][:l-1] == keys[j][:l-1]:
                        _ = keys[i] + (keys[j][l-1],)
                        _id = [self.item2id[k] for k in _]
                        # Obtain co-occurrence counts via multiplication and calculate support
                        support = 1. * sum(np.prod(self.D[:, _id], 1)) / self.total 
                        if support > self.min_support: # Check if it is frequent enough
                            rules[-1][_] = support

        # Traverse each frequent itemset and calculate confidence
        result = {}
        for n,relu in enumerate(rules[1:]): # For all k, traverse k-frequent itemsets
            for r,v in relu.items(): # Traverse all k-frequent itemsets
                for i,_ in enumerate(r): # Traverse all permutations, i.e., is (A,B,C) actually A,B -> C or A,C -> B?
                    x = r[:i] + r[i+1:]
                    confidence = v / rules[n][x] # Confidence of different permutations
                    if confidence > self.min_confidence: # If the confidence is high enough, add to result
                        result[x+(r[i],)] = (confidence, v)

        return sorted(result.items(), key=lambda x: -x[1][0]) # Sort by confidence in descending order