[프로그래머스] 더 맵게 python

문제 링크

https://programmers.co.kr/learn/courses/30/lessons/42626

문제 설명

간단하게 문제를 설명하면, 인자로 넘어오는 scoville 리스트의 모든 요소들이 지정 스코빌 지수인 K보다 크거나 같게 만들기 위해 가장 작은 스코빌 지수와 두번째로 작은 스코빌 지수를 a+b**2 연산을 한다. 그리고 모든 요소들이 K 이상의 값을 갖게 되면 몇 번의 연산을 통해 조건을 만족했는지 cnt를 반환하는 문제이다.

1차 시도

def solution(scoville, K):
    cnt = 0
    scv = 0
    while(1):
        if(len(scoville)<=1):
            return -1
        scoville.sort()
        a = scoville.pop(0)
        b = scoville.pop(0)
        scv = a + b*2
        scoville.append(scv)
        cnt+=1
        for i in scoville:
            if(i<K):
                break
        else:
            return cnt

처음엔 코드를 이렇게 작성했었다. 인자로 받은 리스트를 정렬하고 앞의 두 수를 뽑아내 연산 후 다시 리스트에 집어 넣고 조건을 충족시키는지 확인했다.

정확성 테스트는 통과했으나 효율성 테스트를 통과하지 못했다.
그래서 질문하기 를 둘러보니 heapq를 써야만 효율성 테스트를 통과할 수 있다는 글을 봐서 heapq가 무엇인지 찾아보니 heapq 모듈은 min heap 자료구조를 제공하는데

heapq 모듈 간략 설명

• min heap의 0번째 인덱스엔 가장 작은 값이 위치하고, 원소가 추가되거나 삭제될 때 min heap 내의 모든 원소(k)는 항상 자식 원소들(2k+1, 2k+2) 보다 크기가 작거나 같도록 유지된다.
• heapq 모듈의 함수론 heappush, heappop, heapify 등이 있다.
• 또한 보통 리스트처럼 최소힙을 다룰 수 있다.
• 보통 리스트처럼 인덱스를 통해서 접근 가능한데, 0번째 인덱스에 위치한 값은 항상 최소값이지만, n번째 인덱스에 위치한 값이 n번째로 작은 값이라는 의미는 아니다.

이걸 사용해서 코드를 일부 수정해보았다.

2차 시도

import heapq

def solution(scoville, K):
    heapq.heapify(scoville)
    cnt = 0
    while(1):
        if(len(scoville)<=1):
            return -1
        a = heapq.heappop(scoville)
        b = heapq.heappop(scoville)
        heapq.heappush(scoville, a + b*2)
        cnt+=1
        for i in scoville:
            if(i<K):
                break
        else:
            return cnt

heapq 모듈을 import 해주고 초반에 인자로 받은 scoville 리스트를 heapify를 이용해 힙으로 변환해주었다. 그리고 나머지는 pop과 push 부분에서만 변경해주었다.

이러니 효율성 테스트 통과!
level1의 소수 찾기 문제도 에라토스테네스의 체를 쓰지 않으면 효율성 테스트를 통과할 수 없었는데.. 얘두 heapq를 모르면 효율성 테스트를 통과할 수 없었다.
파이썬 공부를 따로 해줘야겠다.

여기서 내가 살짝 실수했던 부분이 있었는데 힙 자체가 정렬되어있는줄 알고 for문에서 인덱스로 접근했다. 이렇게 하는게 더 효율성이 좋겠군! 하면서 코드를 작성했는데 그렇지 않고 for in을 쓰는게 더 효율성이 좋았다.
이는 위에도 써놨듯이 heap이 인덱스 0번째에 최소값이 있을 뿐, 그게 힙이 정렬되어있다는 것을 의미하는게 아님! 유의하기.

정리

• heapq가 중요한 부분이었다.
• heapq 참조 블로그 - https://www.daleseo.com/python-heapq/