[ps] 치킨 배달

치킨 배달
나의 처음 풀이 : itertools의 combinations를 활용해 살아남을 치킨가게 조합을 구하고, 각 케이스에서 bfs로 각 집당 치킨거리를 구해서 더한다. => 이 값의 최솟값을 최종 결과로 제출.

개선 : python의 generator를 사용. 집에서 각 치킨집까지의 거리를 미리 계산해둔다.

n, m = map(int, input().split())
town = [list(map(int, input().split())) for _ in range(n)]

chicken = []
houses = []

for i in range(n):
    for j in range(n):
        if town[i][j] == 2:
            chicken.append((i, j))
        elif town[i][j] == 1:
            houses.append((i, j))

# 집에서 각 치킨집까지의 거리 미리 계산
distances = [[0] * len(chicken) for _ in range(len(houses))]
for h_idx, (hx, hy) in enumerate(houses):
    for c_idx, (cx, cy) in enumerate(chicken):
        distances[h_idx][c_idx] = abs(hx - cx) + abs(hy - cy)


# 최소 치킨 거리 계산
def calculate_min_chicken_distance(selected_chickens):
    total_distance = 0
    for h_idx in range(len(houses)):
        # 현재 집에서 선택된 치킨집까지의 최소 거리
        min_dist = min(distances[h_idx][c_idx] for c_idx in selected_chickens)
        total_distance += min_dist
    return total_distance


def generate_combinations(chickens, m, idx=0, current_comb=[]):
    if len(current_comb) == m:
        yield current_comb
        return
    
    if idx >= len(chickens):
        return
    
    # 포함하는 경우
    yield from generate_combinations(chickens, m, idx + 1, current_comb + [chickens[idx]])
    # 포함하지 않는 경우
    yield from generate_combinations(chickens, m, idx + 1, current_comb)

result = 1e9
for selected_chickens in generate_combinations(range(len(chicken)), m):
    result = min(result, calculate_min_chicken_distance(selected_chickens))

print(result)