[Python] 프로그래머스 Lv.3 - 야근 지수

요즘 경진대회 준비하며 '프로그래머스(Programmers.co.kr)' 플랫폼을 활용하여 코딩테스트 겸 알고리즘 공부를 하고 있습니다. 이번 글은 Python을 이용하여 프로그래머스의 난이도 Level 3 문제인 '야근 지수' 문제를 풀어보려고 합니다.

프로그래머스

 

문제

문제 전문은 링크를 통해 직접 확인해주시기 바랍니다.

입력되는 값은 정수 배열 works, 정수 n입니다.

1. 입력된 n이 1씩 줄어들면서 works의 원소 중 최대값을 가진 원소에서는 1씩 감소시킵니다.
2. n이 0이 되면 works의 모든 원소를 제곱하여 더합니다

 

제공된 입출력 예시 1번 (works = [4, 3, 3] / n = 4)을 순서대로 보면 다음과 같습니다.

1. works [ 4, 3, 3 ] / n = 4
2. works [ 3, 3, 3 ] / n = 3
3. works [ 2, 3, 3 ] / n = 2 → 같은 값인 원소의 경우 순서대로 처리됨
4. works [ 2, 2, 3 ] / n = 1
5. works [ 2, 2, 2 ] / n = 0 

Result = 2^2 + 2^2 + 2^2 = 12

 

 

첫 번째 시도

첫 시도는 희망적으로 처참했습니다. 분명 입출력 케이스는 모두 정확하게 확인이 됐고, 실채점에서도 정확도는 문제가 없었으나 효율성이 떨어졌습니다. 아무래도 작업을 한 번 수행할 때마다 sort를 통해 모든 원소의 순서를 재배열하는 과정이 반복되기 때문에 시간이 길어진 탓입니다.

def solution(n, works):
    answer = 0
    
    if n >= sum(works): # 1
        return 0
    
    works.sort() # 2
    for _ in range(n): # 3
        works[-1] -= 1 # 3-1
        works.sort() # 3-2
    
    for w in works:
        answer += w ** 2 # 4
        
    return answer

정상적으로 작동하는 코드지만, 효율성이 떨어져 실패함

 

위에서 고안한 알고리즘은 다음과 같습니다.

1. works의 모든 원소를 다 합쳤을 때 n보다 작거나 같으면 계산 후 값이 0으로 동일하므로 조건문 처리
2. works의 모든 원소를 순서대로 정렬 (높은 값이 맨 앞으로)
3. n만큼 반복
3-1. works의 맨 뒤에 있는 원소[-1]의 값을 1 감소시킴
3-2. works의 모든 원소를 순서대로 정렬 (높은 값이 맨 앞으로)
4. 모든 원소를 제곱하여 변수 answer에 합함

 

 

두 번째 시도

같은 정렬이라도 방법에 따라 시간, 복잡도는 천차만별입니다. 두 번째 시도에서는 Heap을 사용했습니다. Heap은 제일 작은 값이 앞으로 옵니다. 별도의 sort 명령어가 없더라도 heap에 들어간 값은 정렬이 되는 셈입니다. 즉 불필요한 sort를 쓰지 않아도 된다는 말입니다.

Python에서 Heap을 사용하기 위해 'heapq' 모듈을 불러왔습니다. 

전체적인 흐름은 첫 번째 시도와 동일합니다. 다만 이번에는 heap에서 값을 불러오고, 다시 넣는 과정이 추가되었다고 보면 됩니다.

import heapq

def solution(n, works):
    answer = 0
    
    if n >= sum(works):
        return 0
    
    works = [-w for w in works] # Heap은 최소값부터 정렬되기 때문에 음수로 바꾸기
    heapq.heapify(works) # HEAP!
    for _ in range(n):
        maximum = heapq.heappop(works) # heap에서 맨 앞에 있는거 (최대값) 가져옴
        heapq.heappush(works, maximum+1) # 1을 더해서 (빼서) 다시 넣고
    
    for w in works:
        answer += w ** 2
        
    return answer

 

 

마치며

분명 쉬운 문제였을 겁니다. 그럼에도 단편적으로만 접근하며 오랜 시간 답을 찾지 못한 저에게 아쉬움이 남습니다. 알고리즘 공부에 더 정진해야 한다는 현실에 마주하게 된 만큼 다양한 문제를 접하고, 해결하는 스킬을 더 늘려야겠다는 다짐을 하며 글을 마칩니다. 

 

오늘도 '응 답 없음' 없는 하루 보내시기 바랍니다.