src/processing.py

@@ -1,8 +1,16 @@
 from preprocessing import parse_file
-from pandas import DataFrame
+from pandas import DataFrame, Series
 from sys import argv
 
 
+def get_furthest_element(element, data):
+    element_df = data.query(f"source == {element} or destination == {element}")
+    furthest_index = element_df["distance"].idxmax()
+    furthest_row = data.iloc[furthest_index]
+    print(furthest_row)
+    return furthest_row
+
+
 def get_first_solution(n, data):
     distance_sum = DataFrame(columns=["point", "distance"])
     for element in range(n):
@@ -16,36 +24,16 @@ def get_first_solution(n, data):
     return furthest_row
 
 
-def get_different_element(original, row):
-    if row.source == original:
-        return row.destination
-    return row.source
-
-
-def get_furthest_element(element, data):
-    element_df = data.query(f"source == {element} or destination == {element}")
-    furthest_index = element_df["distance"].idxmax()
-    furthest_row = data.iloc[furthest_index]
-    furthest_point = get_different_element(original=element, row=furthest_row)
-    furthest_element = {"point": furthest_point, "distance": furthest_row["distance"]}
-    return furthest_element, furthest_index
-
-
 def greedy_algorithm(n, m, data):
     solutions = DataFrame(columns=["point", "distance"])
     first_solution = get_first_solution(n, data)
     solutions = solutions.append(first_solution, ignore_index=True)
     for _ in range(m):
-        last_solution = solutions["point"].tail(n=1)
-        centroid, furthest_index = get_furthest_element(
-            element=int(last_solution), data=data
-        )
-        solutions = solutions.append(dict(centroid), ignore_index=True)
-        data = data.drop(furthest_index)
-    return solutions
+        centroid = solutions.apply(get_furthest_element, 1, data)
+        solutions = solutions.append(centroid)
 
 
-# NOTE In each step, switch to the element that gives the least amount
+# NOTE In each step, switch the element that gives the least amount
 def local_search():
     pass
 
@@ -59,8 +47,7 @@ def main():
     if len(argv) != 2:
         usage(argv)
     n, m, data = parse_file(argv[1])
-    solutions = greedy_algorithm(n, m, data)
-    print(solutions)
+    greedy_algorithm(n, m, data)
 
 
 if __name__ == "__main__":