Counter 사용 [SWEA] D3 S/W 문제해결 기본 1일차 - Flatten

문제 링크

평탄화라고 하여 정해진 숫자 만큼 배열의 가장 큰 수를 하나를 빼고 가장 작은 수 하나를 더해주는 방식으로, 시간이 지날 수록 배열의 요소들이 모두 같은 값을 가지거나 두개의 값만 1이 차이 나도록 해준다.

이 문제를 풀기 위하여 우선 입력받은 배열을 Counter로 바꿔주었다.

from collections import Counter

T = 10
for test_case in range(1, T + 1):
    dumps = int(input())
    blocks = list(map(int,input().split()))
    blocks_counter = Counter(sorted(blocks))
    sorted_blocks = sorted(blocks)
    min_index = sorted_blocks[0]
    max_index = sorted_blocks[-1]

    res = max_index-min_index
    for i in range(dumps):
        if res<=1:
            break
        blocks_counter[min_index]-=1
        blocks_counter[min_index+1]+=1
        blocks_counter[max_index]-=1
        blocks_counter[max_index-1]+=1
        
        if blocks_counter[min_index]==0:
            min_index+=1
        if blocks_counter[max_index]==0:
            max_index-=1
        res = max_index-min_index

    print(f'#{test_case} {res}')

그 다음 카운터의 가장 작은 값을 인덱스로 하는 요소의 개수가 0이 될때까지 계속 빼주며 그만큼 가장 작은 값+1의 수를 1씩 늘려준다.
Counter의 경우 없는 값은 알아서 할당을 새로 해주기 때문에, 편하게 인덱스에 새로운 값을 넣어도 된다.

반대로 가장 큰 값 역시 개수가 0이 될때까지 빼주고, 그만큼 가장 큰값-1의 수를 1씩 늘려준다.