Computer Science/Algorithm

[Python 3] BOJ 26093 고양이 목에 리본 달기

무니화니 2024. 2. 24. 09:58

 

평소에 약한 DP를 연습하기 위해서 해당 문제를 풀었다.

이 문제에서는 고양이들을 일자로 세워두는데, 양 옆에 같은 리본을 둔 고양이를 세워둘 수 없다.

그래서, 조금 나이브하게 접근을 했는데, 만약 가능하면 전 고양이한테 제일 큰 만족도를 얻을 수 있는 리본을 해주려고 한다. 만약, 전 인덱스에 있는 고양이들 중 제일 큰 만족도를 가진 고양이가 해당 리본을 하고 있다면, 두번째로 많은 만족도를 가진 고양이에서 따온다. 

이를 더욱 간결하게 가져오기 위해서, 파이썬의 heapq 라이브러리를 사용했다.

maximum heap를 구현해서 만족도 1등, 2등을 매번 효율적으로 가져올 수 있도록 하였다.

 

 

import heapq
n,k=map(int,input().split())
data=list(list(map(int,input().split())) for _ in range(n))
pastanswer=[[0,-1],[0,-1]]
for i in range(n):
    first,firstindex=pastanswer[0]
    second,secondindex=pastanswer[1]
    q=[]
    for j in range(k):
        if j!=firstindex:
            heapq.heappush(q,[-first-data[i][j],j])
        else:
            heapq.heappush(q,[-second-data[i][j],j])
    first,firstindex=heapq.heappop(q)
    second,secondinex=heapq.heappop(q)
    pastanswer=[[-first,firstindex],[-second,secondindex]]
print(-first)