big speed up

now uses generators and numpy

generate.py

@@ -2,6 +2,9 @@ import argparse
 import math
 import itertools
 import random
+import numpy as np
+import string
+import random
 from pathlib import Path
 from datetime import datetime
 
@@ -12,6 +15,7 @@ def default_filename():
     # the current time in ISO format (ish)
     return datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f") + ".txt"
 
+
 parser = argparse.ArgumentParser()
 parser.add_argument("nodes", type=int)
 parser.add_argument("-o", dest="file", default=default_filename(), help="output file")
@@ -28,35 +32,46 @@ num_nodes = args.nodes
 num_charging = args.charging
 num_connections = args.edges or fully_connected(num_nodes)
 
+if num_connections > fully_connected(num_nodes):
+    print("too many connections (> complete graph)")
+    exit(1)
+
 min_distance = args.min
 max_distance = args.max
 
-alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+alphabet = list(string.ascii_uppercase)
 repeat = 1 + int(math.log(num_nodes, len(alphabet)))
-node_names = list(itertools.islice(itertools.product(alphabet, repeat=repeat), num_nodes))
 
-nodes = {''.join(n): [] for n in node_names} # map node name to list of nodes connected to it
+join_string = lambda x: ("".join(y) for y in x)
+
+node_names = join_string(
+        itertools.islice(
+            itertools.product(alphabet, repeat=repeat),
+            num_nodes
+            )
+        )
 
 # generate paths
-connected = list(itertools.combinations(nodes.keys(), 2))
-random.shuffle(connected)
-connected = connected[:num_connections]
-for a, b in connected:
-    nodes[a].append(b)
-
-charging = list(nodes.keys())
-random.shuffle(charging)
-charging = charging[:num_charging]
+connections = set()
+while len(connections) < num_connections:
+    pair = tuple(random.sample(range(num_nodes), 2))
+    connections.add(pair)
+
+charging = np.zeros(num_nodes, dtype=bool)
+selected_charging = np.random.choice(np.arange(num_nodes), num_charging, replace=False)
+charging[selected_charging] = True
+
 with out.open("w") as f:
     f.write(f"{num_nodes} {num_connections}\n")
     # write nodes and charging status
-    for n in nodes:
-        f.write(f"{n} {int(n in charging)}\n")
+    for i, n in enumerate(node_names):
+        f.write(str(n) + " " + str(int(charging[i])) + "\n")
     # write connections
-    keys = list(nodes.keys())
-    for a, v in nodes.items():
-        for b in v:
+    for a, b in connections:
+        try:
             distance = random.randint(min_distance, max_distance)
-            f.write(f"{keys.index(a)} {keys.index(b)} {distance}\n")
+            f.write(str(a) +  " " + str(b) + " " + str(distance) + "\n")
+        except ValueError:
+            print(f"{a=} {b=}")
 
 print("saved to", out.name)