import numpy as np import itertools def solution(times, times_limit): # find shortest path nodes = len(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 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 break; shortest = new_shortest for i in range(nodes): if shortest[i][0][i] < 0: infinite_time = True break; if infinite_time: return list(range(nodes-2)) shortest = np.reshape(shortest, (nodes, nodes)) #print shortest # calc time for all paths best_plan = [] for order in itertools.permutations(range(1, nodes)): 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 return best_plan 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])