[Python/파이썬] 프로그래머스 야근 지수

프로그래머스

코드

- 틀린 코드!!

def solution(n, works):
    answer = 0
    remain_total = sum(works) - n
    len_w = len(works)
    if remain_total <= 0:
        return 0
    else:
        a = remain_total // len_w
        b = remain_total % len_w
        return (a**2)*(len_w - b) + ((a+1)**2)*b

처음에는 위와 같이 풀었는데 생각해보니 works 배열에 있는 수가 무조건 a보다 크다는 보장이 없어서 이 풀이는 당연히 안된다. 그래서 아래와 같이 풀었다

def solution(n, works):
    answer = 0
    if len(works) == 1:
        return (works[0] - n) ** 2 if (works[0] - n) > 0 else 0
    works = sorted(works, reverse=True)
    idx = 0
    while n > 0:
        works[idx] -= 1
        idx += 1
        n -= 1
        if idx == len(works) or works[idx-1] >= works[idx]:
            idx = 0
    
    if works[0] < 0:
        return 0
    
    for i in works:
        answer += i ** 2
    return answer

works 배열을 내림차순으로 정렬해주고 앞에서부터 배열의 값을 1씩 빼주는데 배열의 끝까지 가거나 방금 1을 빼준 값이 다음 원소의 값보다 크거나 같다면 인덱스를 다시 0으로 초기화시켜준다.

왜 이렇게 해주냐면 배열의 각 원소들의 제곱의 합이 가장 작게 하려면 배열의 원소들 간의 차를 작게 만들어야 하기 때문이다. 예를 들어 남은 작업량이 [2,3,4]인 경우와 [3, 3, 3]인 경우가 있다고 하면 이 두 경우의 남은 총작업량은 같지만 원소들간의 차가 더 작은 [3,3,3]이 문제에서 요구하는 작업량의 제곱값의 합이 [2,3,4] 배열보다 더 작다. 따라서 works 배열에서 큰 값들부터 작업하여 남은 작업량끼리의 차를 줄여주는 코드를 작성한 것이다.