level4/running-with-bunnies/solution.py

@@ -1,52 +1,53 @@
 import numpy as np
-import itertools
 
 def solution(times, times_limit):
-    # find shortest path
-    nodes = len(times)
+    best_times = get_best_times(times)
+    best_plan = get_best_plan(best_times, times_limit)
+    return best_plan
+
+def get_best_times(times):
     t = np.array(times)
-    shortest = np.reshape(t, (nodes, 1, nodes))
-    one_step = np.reshape(t.T, (1, nodes, nodes))
-    #print shortest
-    #print one_step
+    spots = len(times)
+    bunny_count = spots-2
+    best_times = np.reshape(t, (spots, 1, spots))
+    one_step = np.reshape(t.T, (1, spots, spots))
     infinite_time = False
     while not infinite_time:
-        temp = np.reshape(np.amin(shortest + one_step, axis=2), (nodes, 1, nodes))
-        #print temp
-        new_shortest = np.minimum(shortest, temp)
-        #print new_shortest
-        if np.array_equal(shortest, new_shortest):
-            shortest = new_shortest
+        temp = np.reshape(np.amin(best_times + one_step, axis=2), (spots, 1, spots))
+        new_best_times = np.minimum(best_times, temp)
+        if np.array_equal(best_times, new_best_times):
+            best_times = new_best_times
             break;
-        shortest = new_shortest
-        for i in range(nodes):
-            if shortest[i][0][i] < 0:
+        best_times = new_best_times
+        for i in range(spots):
+            if best_times[i][0][i] < 0:
                 infinite_time = True
                 break;
     if infinite_time:
-        return list(range(nodes-2))
-    shortest = np.reshape(shortest, (nodes, nodes))
-    #print shortest
+        return list(range(bunny_count))
+    best_times = np.reshape(best_times, t.shape)
+    return best_times
 
-    # calc time for all paths
+def get_best_plan(best_times, times_limit):
     best_plan = []
-    for order in itertools.permutations(range(1, nodes)):
+    last_position = 0
+    bulkhead = len(best_times)-1
+    bunny_count = len(best_times)-2
+    for plan in get_plans(bunny_count):
         estimated_time = 0
-        order = (0,) + order
-        bunnies = []
-        for i in range(nodes-1):
-            start_node, end_node = order[i], order[i+1]
-            estimated_time += shortest[start_node][end_node]
-            if end_node == nodes-1:
-                break;
-            bunnies += [end_node-1]
-        if estimated_time <= times_limit and len(bunnies) > 0:
-            bunnies.sort()
-            best_plan = get_better_plan(best_plan, bunnies)
-    #print best_plan
-
+        for bunny_id in plan:
+            next_position = bunny_id+1
+            estimated_time += best_times[last_position][next_position]
+            last_position = next_position
+        estimated_time += best_times[last_position][bulkhead]
+        if estimated_time <= times_limit:
+            best_plan = get_better_plan(best_plan, plan)
     return best_plan
 
+# TODO: implement get_plans
+def get_plans(bunny_count):
+    return [[1], [0], [1, 2], [0, 1, 2], [0, 1]]
+
 def get_better_plan(a, b):
     if len(a) > len(b):
         return a