[플래티넘 5] 백준 1102 - 발전소 (파이썬)

https://www.acmicpc.net/problem/1102

풀이

• 발전소가 N개 있고, 어떤 발전소는 작동 중(y), 어떤건 꺼져 있다. (N)
• 작동 중인 발전소가 다른 발전소를 켜는 데 드는 비용이 주어진다. (A[i][j])
• 발전소를 P개 이상 켜기 위해 최소 비용을 구하는 문제이다.
• 켜는 순서와 조합에 따라 총 비용이 달라진다.
• 즉 최적의 순서를 구해야한다.

주요 알고리즘은 비트마스킹 + 메모이제이션을 사용한다.

입력 처리는 다음과 같이 비교적 간단하게 입력 받는다.

N = int(input())
A = [list(map(int, input().split())) for _ in range(N)]
S = input()
P = int(input())

캐시 초기화를 위해서 65536개의 캐시를 생성한다.

cache = [-1] * (2 ** 16)

최대 발전소 수가 16개 이므로 2 ** 16개 즉 65536개의 비트마스크를 가진다.

cache[mask]는 mask 상태에서 P개 이상 켜기 위한 최소 비용을 나타낸다.

dp(mask) 함수에 대해서 설명하면 다음과 같다.

def dp(mask):
    if cache[mask] != -1:
        return cache[mask]

이미 계산된 상태는 재사용 한다.

count = 0
for i in range(N):
    if mask & (2 ** i) != 0:
        count += 1

목표 수 P이상 켜져 있으면 추가 비용이 없다.

ret = 1e9
for i in range(N):
    if mask & (2 ** i) != 0:  # i번 발전소가 켜져 있다면
        for j in range(N):
            if mask & (2 ** j) == 0:  # j번 발전소가 꺼져 있다면
                ret = min(ret, A[i][j] + dp(mask + (2 ** j)))

켜진 발전소 i로 꺼진 발전소 j를 켜는 모든 경우를 고려한다.

j를 켜고 새로운 상태 mask + (2 ** j)로 dp를 진행한다.

cache[mask] = ret
return ret

최소 비용을 저장 및 반환한다.

초기 상태는 다음과 같이 구한다.

init_mask = 0
for i in range(N):
    if S[i] == 'Y':
        init_mask += (2 ** i)

처음 켜져 있는 발전소들을 비트마스크로 표현한다.

answer = dp(init_mask)
if answer == 1e9:
    print(-1)
else:
    print(answer)

결과는 만약 1e9면 P개 이상을 못켰다는 뜻으로 -1을 출력한다.

요약하자면 다음과 같다.

• dp(mask)는 현재 켜져 있는 발전소 상태에서 P개 이상 켜기 위한 최소 비용이다.
• mask는 비트마스킹으로 상태를 표현한다.
• 각 켜진 발전소 i가 꺼진 j를 켜는 경우를 모두 시도한다. (A[i][j])
• 최적 부분 구조와 중복 호출이 있으므로 DP + 메모이제이션을 사용한다.
• 시간복잡도는 O(N^2 * 2 ^ N),, 복잡한데?

코드

# 백준 1102 - 발전소
# 분류 : 다이나믹 프로그래밍

N = int(input())
A = [list(map(int, input().split())) for _ in range(N)]
S = input()
P = int(input())

cache = [-1] * (2 ** 16)
def dp(mask) :
    if cache[mask] != -1 :
        return cache[mask]

    count = 0
    for i in range(N) :
        if mask & (2 ** i) != 0 :
            count += 1
    
    if count >= P :
        cache[mask] = 0
        return 0
    
    ret = 1e9
    for i in range(N) :
        if mask & (2 ** i) != 0 :
            for j in range(N) :
                if mask & (2 ** j) == 0 :
                    ret = min(ret, A[i][j] + dp(mask + (2 ** j)))

    cache[mask] = ret
    return ret

init_mask = 0
for i in range(N) :
    if S[i] == 'Y' :
        init_mask += (2 ** i)

answer = dp(init_mask)
if answer == 1e9 :
    print(-1)
else :
    print(answer)