Generalized python scripts for varying capset length

2 files changed, 41 insertions, 17 deletions

diff --git a/common.py b/common.py
index 411867b..da8fe79 100755
--- a/common.py
+++ b/common.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 from itertools import combinations
+from sys import argv
 
 def increment(dictionary, key):
     dictionary[key] = dictionary.get(key, 0) + 1
@@ -23,7 +24,8 @@ def distribution(counts, capset, num):
     metadist = {}
     for i in range(num_i):
         dists[i] = "".join(map(str, sorted(dists[i])))
-        name = "dist" + str(num) + ":" + dists[i]
+        name = "dist" + str(num) + ":" \
+             + str(len(capset))  + ":" + dists[i]
         increment(counts, name)
         increment(metadist, name)
     if counts["premeta" + str(num)]:
@@ -51,8 +53,15 @@ def metadistribution(counts, num):
             del counts[key]
 
 def main():
+    if len(argv) == 1:
+        infile = "capset.out"
+    elif len(argv) == 2:
+        infile = argv[1]
+    else:
+        print("Usage: common.py [INFILE]")
+        exit()
     capsets = []
-    capset_file = open("capset.out", "r")
+    capset_file = open(infile, "r")
     for line in capset_file:
         capsets.append(line[1:-2].split(" "))
     capset_file.close()
diff --git a/connected.py b/connected.py
index ee3ae70..9c173fe 100755
--- a/connected.py
+++ b/connected.py
@@ -1,9 +1,16 @@
 #!/usr/bin/env python3
 from sys import argv
-if len(argv) == 1: NUM = 1
-if len(argv) == 2: NUM = int(argv[1])
+if len(argv) <= 1: NUM = 1
+if len(argv) >= 2: NUM = int(argv[1])
+
+if len(argv) <= 2: KEY = None
+if len(argv) >= 3: KEY = int(argv[2])
 
 def is_connected(capset0, capset1):
+    """Decides whether the two capsets are connected.
+    
+    In this case, capsets are connected if you can get from
+    one to the other by changing NUM cards or fewer."""
     count = 0
     for card in capset0:
         if card in capset1:
@@ -12,6 +19,7 @@ def is_connected(capset0, capset1):
     return count <= NUM
 
 def next_connected(disconnected):
+    """Finds the next partition of connected capsets."""
     newly_connected = [disconnected.pop()]
     connected = []
     while newly_connected:
@@ -25,22 +33,29 @@ def next_connected(disconnected):
     return connected
 
 def main():
-    capsets = []
+    capsets = {}
     capset_file = open("capset.out", "r")
     for line in capset_file:
-        capsets.append(tuple(line[1:-2].split(" ")))
+        capset = tuple(line[1:-2].split(" "))
+        if len(capset) in capsets:
+            capsets[len(capset)].append(capset)
+        else:
+            capsets[len(capset)] = [capset]
     capset_file.close()
 
-    connecteds = []
-    while capsets:
-        connecteds.append(next_connected(capsets))
-    #for connected in connecteds:
-    #    print("\nThe following are connected:")
-    #    for capset in connected:
-    #        print(capset)
-    print("\nThere are", len(connecteds), "partitions.")
-    lengths = list(map(len, connecteds))
-    for length in sorted(set(lengths)):
-        print("There are", lengths.count(length), "with length", length)
+    connecteds = {}
+    if KEY == None:
+        keys = sorted(capsets.keys())
+    else:
+        keys = [KEY]
+    for key in keys:
+        connecteds[key] = []
+    for key in keys:
+        while capsets[key]:
+            connecteds[key].append(next_connected(capsets[key]))
+        print("\nFor length {} there are {} partitions.".format(key, len(connecteds[key])))
+        lengths = list(map(len, connecteds[key]))
+        for length in sorted(set(lengths)):
+            print("There are {} capsets with length {}.".format(lengths.count(length), length))
 
 if __name__ == "__main__": main()