Implement contingency table

pysie/dsl/set.py

@@ -41,8 +41,8 @@ def _put(self, x, key, value, d):
         elif compared > 0:
             x.right = self._put(x.right, key, value, d)
         else:
-            if len(key)-1 > d:
-                x.mid = self._put(x.mid, key, value, d+1)
+            if len(key) - 1 > d:
+                x.mid = self._put(x.mid, key, value, d + 1)
             else:
                 if x.value is None:
                     self.N += 1
@@ -66,7 +66,7 @@ def _get(self, x, key, d):
             return self._get(x.right, key, d)
         else:
             if len(key) - 1 > d:
-                return self._get(x.mid, key, d+1)
+                return self._get(x.mid, key, d + 1)
             else:
                 return x
 
@@ -83,8 +83,8 @@ def _delete(self, x, key, d):
         elif compared > 0:
             x.right = self._delete(x.right, key, d)
         else:
-            if len(key)-1 > d:
-                x.mid = self._delete(x.mid, key, d+1)
+            if len(key) - 1 > d:
+                x.mid = self._delete(x.mid, key, d + 1)
             else:
                 self.N -= 1
                 x = None
@@ -102,11 +102,41 @@ def size(self):
     def is_empty(self):
         return self.N == 0
 
+    def keys(self):
+        queue = []
+        self.collect(self.root, '', queue)
+        return queue
+
+    def values(self):
+        queue = []
+        self.collect_values(self.root, queue)
+
+    def collect(self, x, prefix, queue):
+        if x is None:
+            return
+        if x.value is not None:
+            queue.append(prefix + chr(x.key))
+        self.collect(x.left, prefix, queue)
+        self.collect(x.mid, prefix + chr(x.key), queue)
+        self.collect(x.right, prefix, queue)
+
+    def collect_values(self, x, queue):
+        if x is None:
+            return
+        if x.value is not None:
+            queue.append(x.value)
+
+        self.collect_values(x.left, queue)
+        self.collect_values(x.mid, queue)
+        self.collect_values(x.right, queue)
 
-class TernarySearchSet(TernarySearchTrie):
 
+class TernarySearchSet(TernarySearchTrie):
     def add(self, key):
         self.put(key, 0)
 
     def contains(self, key):
         return self.contains_key(key)
+
+    def to_array(self):
+        return self.keys()

pysie/dsl/variable_independence_testing.py

@@ -14,12 +14,37 @@ def __init__(self):
     def set_cell(self, row_name, column_name, value):
         key = self.make_key(row_name, column_name)
         self.values.put(key, value)
+        self.rows.add(row_name)
+        self.columns.add(column_name)
 
     def get_cell(self, row_name, column_name):
         key = self.make_key(row_name, column_name)
+        if not self.values.contains_key(key):
+            return 0
         return self.values.get(key)
 
     def make_key(self, row_name, column_name):
         return row_name + '-' + column_name
 
+    def get_row_total(self, row_name):
+        column_names = self.columns.to_array()
+        result = 0
+        for x in column_names:
+            result += self.get_cell(row_name, x)
+        return result
+
+    def get_column_total(self, column_name):
+        row_names = self.rows.to_array()
+        result = 0
+        for x in row_names:
+            result += self.get_cell(x, column_name)
+        return result
+
+    def get_total(self):
+        values = self.values.values()
+        result = 0
+        for val in values:
+            result += val
+        return result
+
 

tests/dsl/set_unit_test.py

@@ -23,5 +23,8 @@ def test_map(self):
         self.assertFalse(trie.contains_key('hi'))
         self.assertEqual(trie.size(), 101)
 
+        keys = trie.keys()
+        self.assertEqual(len(keys), 101)
+
 if __name__ == '__main__':
     unittest.main()