import numpy as np import copy def solution(times, times_limit): spots = len(times) bulkhead_position = spots-1 bunnies = spots-2 best_times = get_best_times(times, spots, bunnies) if best_times is None: return list(range(bunnies)) best_plan = get_best_plan(best_times, times_limit, bunnies, bulkhead_position) return best_plan def get_best_times(times, spots, bunnies): t = np.array(times) 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(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; 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 None best_times = np.reshape(best_times, t.shape) return best_times def get_best_plan(best_times, times_limit, bunnies, bulkhead_position, start_position = 0): best_plan = [] for plan in get_plans([], list(range(bunnies)) + [None]): if get_better_plan(best_plan, plan) == best_plan: continue estimated_time = 0 last_position = start_position 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_position] if estimated_time <= times_limit: best_plan = get_better_plan(best_plan, plan) best_plan.sort() return best_plan def get_plans(saved, remain): for id in remain: remain_bunnies = copy.deepcopy(remain) remain_bunnies.remove(id) if id == None: yield saved for id in get_plans(saved + [id], remain_bunnies): yield id def get_better_plan(a, b): if len(a) > len(b): return a if len(a) < len(b): return b if a == b: return a for i in range(len(a)): if a[i] == b[i]: continue if a[i] < b[i]: return a if a[i] > b[i]: return b return a tests = [ [[[0, 2, 2, 2, -1], [9, 0, 2, 2, -1], [9, 3, 0, 2, -1], [9, 3, 2, 0, -1], [9, 3, 2, 2, 0]], 1, [1, 2]], [[[0, 1, 1, 1, 1], [1, 0, 1, 1, 1], [1, 1, 0, 1, 1], [1, 1, 1, 0, 1], [1, 1, 1, 1, 0]], 3, [0, 1]], [[[0, 2, 2, 2, -1], [9, 0, 2, 2, -1], [9, 3, 0, 2, -1], [9, 3, 2, 0, -1], [9, 3, 2, 2, 0]], -1, []], [[[0, 2, 2, 2, -1], [9, 0, 2, 2, -1], [9, 3, 0, 2, -1], [9, 3, 2, 0, -1], [9, 3, 2, 2, 0]], 0, [1]], [[[0, 2, 2, 2, -1], [9, 0, 2, 2, -1], [9, 3, 0, 2, -1], [9, 3, 2, 0, -1], [9, 3, 2, 2, 0]], 2, [0, 1]], [[[0, 2, 2, 2, -1], [9, 0, 2, 2, -1], [9, 3, 0, 2, -1], [9, 3, 2, 0, -1], [9, 3, 2, 2, 0]], 3, [0, 1, 2]], ] for test in tests: result = solution(test[0], test[1]) print (test[0], test[1], result == test[2], result, test[2])