Fix a bug in support calculation

dzzh · enaeseth · commit 9ba1f2d041c5 · 2013-06-19T09:22:46.000-07:00
Fixes enaeseth#2.
diff --git a/fp_growth.py b/fp_growth.py
@@ -232,11 +232,9 @@ def conditional_tree_from_paths(paths, minimum_support):
 
     # 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
+        count = path[-1].count
+        for node in reversed(path[:-1]):
+            node._count += count
 
     # Eliminate the nodes for any items that are no longer frequent.
     for item in items:
@@ -403,7 +401,7 @@ def __repr__(self):
 
     f = open(args[0])
     try:
-        for itemset in find_frequent_itemsets(csv.reader(f), options.minsup):
-            print '{' + ', '.join(itemset) + '}'
+        for itemset, support in find_frequent_itemsets(csv.reader(f), options.minsup, True):
+            print '{' + ', '.join(itemset) + '} ' + str(support)
     finally:
         f.close()
diff --git a/test.py b/test.py
@@ -205,6 +205,19 @@ def testPruning(self):
         root.child('a', 2).leaf()
         self.assertEqual(1, len(root.node.children))
 
+    def testSupport(self):
+        """
+        Issue #2: incorrect support calculation for shared paths
+        """
+        self.tree.add('abcd')
+        self.tree.add('abd')
+        paths = list(self.tree.prefix_paths('d'))
+        ct = fp_growth.conditional_tree_from_paths(paths, 1)
+        root = NodeTester(self,ct.root)
+        a = root.child('a',2)
+        b = a.child('b',2)
+        c = b.child('c',1)
+
 class FrequentSetTests(unittest.TestCase):
     def testDuplicate(self):
         raw = '25,52,274;71;71,274;52;25,52;274,71'
