Added a function to generate a conditional tree from a list of prefix paths.

Eric Naeseth · Eric Naeseth · commit 300633afd15e · 2009-03-10T03:27:19.000-05:00
diff --git a/fp_growth.py b/fp_growth.py
@@ -88,7 +88,53 @@ def collect_path(node):
             return path
             
         return (collect_path(node) for node in self.nodes(item) if node.leaf)
+
+def conditional_tree_from_paths(paths, minimum_support):
+    """Builds 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]
+        
+        point = tree.root
+        for node in path:
+            next_point = point.search(node.item)
+            if not next_point:
+                items.add(node.item)
+                count = node.count if node.leaf else 0
+                next_point = FPNode(tree, node.item, count)
+                point.add(next_point)
+                tree._update_route(next_point)
+            point = next_point
+                
+    # Calculate the counts of the non-leaf nodes.
+    for path in tree.prefix_paths(condition_item):
+        count = None
+        for node in reversed(path):
+            if count is not None:
+                node._count += count
+            count = node.count
+    
+    # Eliminate the nodes for any items that are no longer frequent.
+    for item in items:
+        support = sum(n.count for n in tree.nodes(item))
+        if support < minimum_support:
+            # Doesn't make the cut anymore
+            for node in tree.nodes(item):
+                node.parent.remove(node)
+                
+    # Finally, remove the nodes corresponding to the item for which this
+    # conditional tree was generated.
+    for node in tree.nodes(condition_item):
+        node.parent.remove(item)
         
+    return tree
     
 class FPNode(object):
     """A node in an FP tree."""
