92 lines
3.2 KiB
Python

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])