32 lines
1.2 KiB
Python
32 lines
1.2 KiB
Python
import numpy as np
|
|
|
|
def solution(entrances, exits, path):
|
|
p = np.array(path)
|
|
minimum_width = -1
|
|
all_rooms = set(range(len(p)))
|
|
for rooms, others in generate_rooms(all_rooms, set(entrances), set(exits), p):
|
|
width = calculate_cut_width(rooms, others, p)
|
|
if width < minimum_width or minimum_width == -1:
|
|
minimum_width = width
|
|
return minimum_width
|
|
|
|
def generate_rooms(all_rooms, prev_rooms, exits, path):
|
|
yield prev_rooms, all_rooms - set(prev_rooms)
|
|
for room, width in enumerate(sum(path[list(prev_rooms)])):
|
|
if room not in prev_rooms and room not in exits and width != 0:
|
|
next_rooms = prev_rooms | set([room])
|
|
for result in generate_rooms(all_rooms, next_rooms, exits, path):
|
|
yield result
|
|
|
|
def calculate_cut_width(rooms, others, p):
|
|
return np.sum(p[list(rooms)][:, list(others)])
|
|
|
|
tests = [
|
|
[[[0], [3], [[0, 7, 0, 0], [0, 0, 6, 0], [0, 0, 0, 8], [9, 0, 0, 0]]], 6],
|
|
[[[0, 1], [4, 5], [[0, 0, 4, 6, 0, 0], [0, 0, 5, 2, 0, 0], [0, 0, 0, 0, 4, 4], [0, 0, 0, 0, 6, 6], [0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0]]], 16],
|
|
]
|
|
|
|
for test in tests:
|
|
result = solution(*test[0])
|
|
print(test[0], result == test[1], result, test[1])
|