1 files changed, 74 insertions, 40 deletions

diff --git a/ass2/q1/kruskals.py b/ass2/q1/kruskals.py
index 4b8da4d..f3c0703 100644
--- a/ass2/q1/kruskals.py
+++ b/ass2/q1/kruskals.py
@@ -1,27 +1,39 @@
-from functools import reduce
+# Alexander Occhipinti - 29994705
+
+import sys
 
 
 class Edge:
+    """
+    The class which represents an edge in a graph
+    """
     def __init__(self, start, end, weight):
         self.start = start
         self.end = end
         self.weight = weight
 
-    def __str__(self):
-        return f"{self.start} -({self.weight})-> {self.end}"
-
     def __repr__(self):
-        return str(self)
+        return f"{self.start} {self.end} {self.weight}"
+
+    def __str__(self):
+        return repr(self) + "\n"  # The form output to file
 
-    def __int__(self):
+    def __len__(self):
         return self.weight
 
+
 def create_edge_from_line(line):
+    """
+    Creates an Edge object by parsing a line of the format specified in the assignment brief
+    """
     data = map(int, line.split(" "))
     return Edge(*data)
 
 
 def read_edges(filename):
+    """
+    Reads in a file and creates a list of Edge objects by parsing the format specified in the assignment brief
+    """
     edges = []
     with open(filename, "r") as file:
         lines = file.read().split("\n")
@@ -30,42 +42,44 @@ def read_edges(filename):
     return edges
 
 
-def count_vertices(edges):
-    v = -1
-    for edge in edges:
-        v = max(v, edge.start, edge.end)
-    return v + 1
-
-
 def sort_edges(edges):
-    return sorted(edges, key=lambda e: e.weight)  # Timsort, O(e log e)
+    """
+    Returns a sorted (from smallest to biggest edge) version of a list of edges
+    Complexity is O(e log e) where e = num of edges due to timSort.
+    """
+    return sorted(edges, key=len)
 
 
 def find(groups, vertex):
-    # return vertex if groups[vertex] < 0 else find(groups, groups[vertex])
-    if groups[vertex] < 0:
-        return vertex
-    else:
-        return find(groups, groups[vertex])
+    """
+    Finds the root of the vertex's group according to a union-find data structure
+    """
+    while groups[vertex] >= 0:
+        vertex = groups[vertex]
+    return vertex
 
 
 def union(groups, u, v):
-    root_u = find(groups, u)
-    root_v = find(groups, v)
+    """
+    Unions two vertices together according to a union-find data structure
+    """
+    root_u, root_v = find(groups, u), find(groups, v)
+    height_u, height_v = -groups[root_u], -groups[root_v]
     if root_u == root_v:
         return
-    height_u = -groups[root_u]
-    height_v = -groups[root_v]
     if height_u > height_v:
         groups[root_v] = root_u
     elif height_v > height_u:
         groups[root_u] = root_v
     else:
         groups[root_u] = root_v
-        groups[root_v] = -(height_v+1)
+        groups[root_v] = -(height_v + 1)
 
 
 def is_spanning(min_span_tree, v):
+    """
+    Returns if a given graph touches all vertices
+    """
     exist = [False for _ in range(v)]
     for edge in min_span_tree:
         exist[edge.start] = True
@@ -73,25 +87,45 @@ def is_spanning(min_span_tree, v):
     return all(exist)
 
 
-def main():
-
-    # print(sum(map(int, [2, Edge(0, 1, 5)])))
-
-    edges = read_edges("edges.txt")
-    print(edges)
-    print(sort_edges(edges))
-    v = count_vertices(edges)
-    groups = [-1 for _ in range(v)]
+def kruskals(edges, v):
+    """
+    Generates a minimum spanning tree for an undirected, weighted, connected graph using Kruskal's algorithm
+    """
+    groups = [-1 for _ in range(v)]  # Init tree structure for union-by-height
     min_span_tree = []
-    for edge in sort_edges(edges):
-        if find(groups, edge.start) != find(groups, edge.end):
+    total_weight = 0
+    for edge in sort_edges(edges):  # Start from the smallest edge and go up
+        if find(groups, edge.start) != find(groups, edge.end):  # Add to span tree if edges are from different groups
             union(groups, edge.start, edge.end)
             min_span_tree.append(edge)
+            total_weight += edge.weight
+    return total_weight, min_span_tree
 
-    weight_total = sum(map(int, min_span_tree))
-    print(weight_total)
-    print("\n".join(map(str, min_span_tree)))
-    assert is_spanning(min_span_tree, v)
 
+def write_output(total_weight, min_span_tree, filename):
+    """
+    Writes the required data to a given file in the format required by the assignment
+    """
+    with open(filename, "w") as file:
+        file.write(f"{total_weight}\n")
+        file.writelines(map(str, min_span_tree))
 
-main()
+
+def parse_args(argv):
+    if len(sys.argv) != 3:
+        raise IndexError(f"Incorrect number of arguments provided. Expected 3, got {len(argv)}.")
+    return int(argv[1]), argv[2]
+
+
+def main():
+    """
+    The main function of the program.
+    """
+    v, filename = parse_args(sys.argv)
+    edges = read_edges(filename)
+    total_weight, min_span_tree = kruskals(edges, v)
+    write_output(total_weight, min_span_tree, "output_kruskals.txt")
+    print("Done.")
+
+if __name__ == "__main__":
+    main()