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