关联规则挖掘之FP-growth算法实现

# encoding: utf-8 """ A Python implementation of the FP-growth algorithm. Basic usage of the module is very simple: > from fp_growth import find_frequent_itemsets > find_frequent_itemsets(transactions, minimum_support) """ from collections import defaultdict, namedtuple from itertools import imap __author__ = 'Eric Naeseth <eric@naeseth.com>' __copyright__ = 'Copyright © 2009 Eric Naeseth' __license__ = 'MIT License' def find_frequent_itemsets(transactions, minimum_support, include_support=False): """ Find frequent itemsets in the given transactions using FP-growth. This function returns a generator instead of an eagerly-populated list of items. The `transactions` parameter can be any iterable of iterables of items. `minimum_support` should be an integer specifying the minimum number of occurrences of an itemset for it to be accepted. Each item must be hashable (i.e., it must be valid as a member of a dictionary or a set). If `include_support` is true, yield (itemset, support) pairs instead of just the itemsets. """ items = defaultdict(lambda: 0) # mapping from items to their supports # Load the passed-in transactions and count the support that individual # items have. for transaction in transactions: for item in transaction: items[item] += 1 # Remove infrequent items from the item support dictionary. items = dict((item, support) for item, support in items.iteritems() if support >= minimum_support) # Build our FP-tree. Before any transactions can be added to the tree, they # must be stripped of infrequent items and their surviving items must be # sorted in decreasing order of frequency. def clean_transaction(transaction): transaction = filter(lambda v: v in items, transaction) transaction.sort(key=lambda v: items[v], reverse=True) return transaction master = FPTree() for transaction in imap(clean_transaction, transactions): master.add(transaction) def find_with_suffix(tree, suffix): for item, nodes in tree.items(): support = sum(n.count for n in nodes) if support >= minimum_support and item not in suffix: # New winner! found_set = [item] + suffix yield (found_set, support) if include_support else found_set # Build a conditional tree and recursively search for frequent # itemsets within it. cond_tree = conditional_tree_from_paths(tree.prefix_paths(item)) for s in find_with_suffix(cond_tree, found_set): yield s # pass along the good news to our caller # Search for frequent itemsets, and yield the results we find. for itemset in find_with_suffix(master, []): yield itemset class FPTree(object): """ An FP tree. This object may only store transaction items that are hashable (i.e., all items must be valid as dictionary keys or set members). """ Route = namedtuple('Route', 'head tail') def __init__(self): # The root node of the tree. self._root = FPNode(self, None, None) # A dictionary mapping items to the head and tail of a path of # "neighbors" that will hit every node containing that item. self._routes = {} @property def root(self): """The root node of the tree.""" return self._root def add(self, transaction): """Add a transaction to the tree.""" point = self._root for item in transaction: next_point = point.search(item) if next_point: # There is already a node in this tree for the current # transaction item; reuse it. next_point.increment() else: # Create a new point and add it as a child of the point we're # currently looking at. next_point = FPNode(self, item) point.add(next_point) # Update the route of nodes that contain this item to include # our new node. self._update_route(next_point) point = next_point def _update_route(self, point): """Add the given node to the route through all nodes for its item.""" assert self is point.tree try: route = self._routes[point.item] route[1].neighbor = point # route[1] is the tail self._routes[point.item] = self.Route(route[0], point) except KeyError: # First node for this item; start a new route. self._routes[point.item] = self.Route(point, point) def items(self): """ Generate one 2-tuples for each item represented in the tree. The first element of the tuple is the item itself, and the second element is a generator that will yield the nodes in the tree that belong to the item. """ for item in self._routes.iterkeys(): yield (item, self.nodes(item)) def nodes(self, item): """ Generate the sequence of nodes that contain the given item. """ try: node = self._routes[item][0] except KeyError: return while node: yield node node = node.neighbor def prefix_paths(self, item): """Generate the prefix paths that end with the given item.""" def collect_path(node): path = [] while node and not node.root: path.append(node) node = node.parent path.reverse() return path return (collect_path(node) for node in self.nodes(item)) def inspect(self): print 'Tree:' self.root.inspect(1) print print 'Routes:' for item, nodes in self.items(): print ' %r' % item for node in nodes: print ' %r' % node def conditional_tree_from_paths(paths): """Build a conditional FP-tree from the given prefix paths.""" tree = FPTree() condition_item = None items = set() # Import the nodes in the paths into the new tree. Only the counts of the # leaf notes matter; the remaining counts will be reconstructed from the # leaf counts. for path in paths: if condition_item is None: condition_item = path[-1].item point = tree.root for node in path: next_point = point.search(node.item) if not next_point: # Add a new node to the tree. items.add(node.item) count = node.count if node.item == condition_item else 0 next_point = FPNode(tree, node.item, count) point.add(next_point) tree._update_route(next_point) point = next_point assert condition_item is not None # Calculate the counts of the non-leaf nodes. for path in tree.prefix_paths(condition_item): count = path[-1].count for node in reversed(path[:-1]): node._count += count return tree class FPNode(object): """A node in an FP tree.""" def __init__(self, tree, item, count=1): self._tree = tree self._item = item self._count = count self._parent = None self._children = {} self._neighbor = None def add(self, child): """Add the given FPNode `child` as a child of this node.""" if not isinstance(child, FPNode): raise TypeError("Can only add other FPNodes as children") if not child.item in self._children: self._children[child.item] = child child.parent = self def search(self, item): """ Check whether this node contains a child node for the given item. If so, that node is returned; otherwi