SWEA 1208. Flatten (파이썬)

 

문제
한 쪽 벽면에 다음과 같이 노란색 상자들이 쌓여 있다.
높은 곳의 상자를 낮은 곳에 옮기는 방식으로 최고점과 최저점의 간격을 줄이는 작업을 평탄화라고 한다.
평탄화를 모두 수행하고 나면, 가장 높은 곳과 가장 낮은 곳의 차이가 최대 1 이내가 된다.
평탄화 작업을 위해서 상자를 옮기는 작업 횟수에 제한이 걸려있을 때, 제한된 횟수만큼 옮기는 작업을 한 후 최고점과 최저점의 차이를 반환하는 프로그램을 작성하시오.
(출제 사이트로 가서 이미지 설명을 참고하세요)

제약 사항
가로 길이는 항상 100으로 주어진다.
모든 위치에서 상자의 높이는 1이상 100이하로 주어진다.
덤프 횟수는 1이상 1000이하로 주어진다.
주어진 덤프 횟수 이내에 평탄화가 완료되면 더 이상 덤프를 수행할 수 없으므로 그 때의 최고점과 최저점의 높이 차를 반환한다 (주어진 data에 따라 0 또는 1이 된다).

입력
총 10개의 테스트 케이스가 주어지며, 각 테스트 케이스의 첫 번째 줄에는 덤프 횟수가 주어진다. 그리고 다음 줄에 각 상자의 높이값이 주어진다.

출력
#부호와 함께 테스트 케이스의 번호를 출력하고, 공백 문자 후 테스트 케이스의 최고점과 최저점의 높이 차를 출력한다.

문제에 대한 모든 권리는 SW Expert Academy에게 있습니다.

 

 

 

 아이디어 

카운팅 정렬을 사용해 높이별 상자의 개수를 정렬하고, 양 끝에서 하나씩 줄여나가자

 

 코드 

for tc in range(1, 11):
    dump_chance = int(input())
    boxes = list(map(int, input().split()))
    heights = [0] * 100

    # counting sort
    for box in boxes:
        for idx in range(100):
            if box == idx+1:
                heights[idx] += 1

    min_h = max_h = 0
    # 최소 인덱스는 앞에서부터 순서대로 탐색
    for idx in range(100):
        if heights[idx] > 0:
            min_h = idx
            break
    # 최대 인덱스는 뒤에서부터 거꾸로 탐색
    for idx in range(99, -1, -1):
        if heights[idx] > 0:
            max_h = idx
            break

    # dump하면서 가장 높은 박수의 개수 줄이고, 가장 낮은 박스의 개수도 줄이자
    for dump in range(dump_chance):
        # 가장 높은 곳과 가장 낮은 곳의 차이가 1이하라면
        # 더이상 덤프하지 않아도 되므로 break
        if max_h - min_h <= 1:
            break
        
        # 2 이상의 차이가 난다면 평탄화 작업 ㄱ
        else:
            # 가장 높은 곳에서 상자를 뺀다 == 가장 높은 곳-1개인 곳이 한 개 생긴다
            heights[max_h] -= 1
            heights[max_h - 1] += 1
            # 가장 높은 곳 상자를 다 썼다면
            if heights[max_h] == 0:
                # 인덱스를 옮겨서 한칸 앞으로(왼쪽 방향으로)
                max_h -= 1

            # 가장 낮은 곳에 상자를 더한다 == 가장 낮은 곳+1개인 곳이 한 개 생긴다
            heights[min_h] -= 1
            heights[min_h + 1] += 1
            # 가장 낮은 곳에 상자가 없다면(다 하나씩 많아져서)
            if heights[min_h] == 0:
                # 인덱스를 한칸 뒤로 옮긴다(오른쪽 방향으로)
                min_h += 1

    print(f'#{tc} {max_h - min_h}')