[Python/Codetree] 병원 거리 최소화하기

https://www.codetree.ai/training-field/frequent-problems/problems/min-of-hospital-distance/description?page=1&pageSize=20&order=tier 

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

문제

사람, 병원 혹은 빈 칸으로 이루어져 있는 n×n 크기의 도시가 주어집니다. 이 때, 각 사람의 병원 거리는 가장 가까운 병원까지의 거리를 의미하며, 두 격자 (x1​,y1​), (x2​,y2​) 사이의 거리는 ∣x2​−x1​∣+∣y2​−y1​∣으로 정의됩니다. 운영상의 이유로 m개의 병원만을 남겨두고 나머지를 폐업해야 할 때 남은 m개의 병원에 대한 각 사람들의 병원 거리의 총 합이 최소가 되도록 하려고 합니다.

예를 들어 아래의 그림과 같은 상황에서 하나의 병원만 남겨두려고 할 때

아래의 그림과 같이 병원을 고를 경우 병원 거리의 총 합이 2 이지만

다른 병원을 고를 경우 병원 거리의 총 합은 4가 됩니다.

m개의 병원을 잘 골라 병원거리의 총 합을 최소화하는 프로그램을 작성해보세요.

입력 형식

첫 번째 줄에는 n과 m이 공백을 사이에 두고 주어지고, 두 번째 줄부터 (n+1)번째 줄까지는 각 행에 빈 칸인 경우 0, 사람인 경우 1, 병원인 경우 2으로 입력이 공백을 사이에 두고 주어집니다. m 개의 병원을 고르는게 불가능한 입력은 주어지지 않는다고 가정해도 좋습니다.

• 2≤n≤50
• 1≤m≤13
• m≤ 병원의 수 ≤13
• 사람의 수 ≤2∗n

출력 형식

병원 m개를 남겼을 때 가능한 각 사람들의 병원 거리 총 합 중 최솟값을 출력하세요

풀이

Code

# 1. 사람, 병원 리스트 생성
people, all_hospital, hospital = [], [], []
L = 0
# 2. 최솟값 변수 생성
ans = float('inf')

# 3. 백트래킹 함수 생성
def backtracking(idx, count) :
    global ans, L
    # 3-1. 종료 조건 설정
    if count == m :
        # 최솟값 업데이트
        distance = 0
        for p_x, p_y in people :
            distance += min([abs(h_x - p_x) + abs(h_y - p_y) for h_x, h_y in hospital])
        ans = min(ans, distance)
        return

    # 3-2. 
    for i in range(idx, L) :
        # 리스트에 해당 병원 인덱스 삽입
        hospital.append(all_hospital[i])
        # 백트래킹
        backtracking(i+1, count+1)
        # 리스트에 해당 병원 인덱스 제거
        hospital.pop()

# 4. 솔루션 함수
def solution(n, m, graph) :
    global ans, L
    # 4-1.
    for i in range(n) :
        for j in range(n) :
            # 사람, 병원 정보 입력
            if graph[i][j] == 1 : people.append((i, j))
            elif graph[i][j] == 2 : all_hospital.append((i, j))
    L = len(all_hospital)
    # 4-2. 백트래킹 실행
    backtracking(0, 0)
    # 4-3. 최솟값 출력
    print(ans)
 
if __name__ == "__main__" :
    n, m = map(int, input().split())
    graph = [list(map(int, input().split())) for _ in range(n)]
    solution(n, m, graph)