[BOJ3661] 생일 선물 - Python

문제

풀이

정렬
금액 낮은 순
금액 같을 시 index 큰 순

H = 공평히 낼 수 있는 금액
W = 금액을 감당할 수 있는 사람

p 에서 H * W 만큼 삭감

H를 줄인 후 다시 연산
(H의 인원이 모두 공평하게 돈을 낼 수 있을 만큼 왼쪽부터 배제)

p가 0이 될 때 까지 반복

if p=3

코드

import sys
ip = lambda:map(int, sys.stdin.readline().split())
T, *_ = ip()
for _ in range(T):
    p, n = ip()
    f = ip()
    # [인덱스, 최대금액, 최대금액(연산용)]으로 저장
    arr = [[idx, n, n] for idx, n in enumerate(f)]
    # 금액 낮은 순 정렬
    # 같은 금액 당 높은 인덱스 순 정렬 
    arr.sort(key=lambda x: (x[1], -x[0]))

    # 쌓이는 H길이
    s = 0
    for i in range(n):
        # 요소당 s만큼 빼기
        arr[i][2] -= s
        if arr[i][2] == 0:
            continue
        # W가 남은 금액보다 큰 경우
        if n-i > p:
            continue
        # 뺄 수 있는 만큼 빼기
        d = min(arr[i][2], p // (n-i))
        if d == 0:
            continue
        arr[i][2] -= d
        p -= d * (n-i)

        # L추가
        s += d
    if p > 0:
        print("IMPOSSIBLE")
        continue

    arr.sort(key=lambda x:x[0])

    for i in arr:
        print(i[1] - i[2], end=" ")
    print()