1 files changed, 30 insertions, 15 deletions
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()