Adding FPTree#nodes(item), which generates all of the nodes that contain the given item.

Eric Naeseth · Eric Naeseth · commit 2c14ab415ed6 · 2009-03-10T02:18:24.000-05:00
diff --git a/fp_growth.py b/fp_growth.py
@@ -53,6 +53,20 @@ def add(self, transaction):
                     self._paths[item] = [next_point, next_point]
                     
             point = next_point
+            
+    def nodes(self, item):
+        """
+        Generates the sequence of nodes that contain the given item.
+        """
+        
+        try:
+            node = self._paths[item][0]
+        except KeyError:
+            return
+            
+        while node:
+            yield node
+            node = node.neighbor
         
     
 class FPNode(object):
