AtCoder: D - Integer Cards - 파이썬(Python)

https://atcoder.jp/contests/abc127/tasks/abc127_d

D - Integer Cards

<문제 선지>

$N$장의 카드가 있다. $i$번째 카드에 정수 $A_i$가 쓰여 있다.
이 순서로 $j=1,2,...,M$마다 다음 연산을 한 번 수행한다.

연산: 최대 $B_j$장의 카드를 선택한다(0장 이상). 선택한 각 카드에 쓰여진 정수를 $C_j$로 변경한다.

$M$ 연산 후, $N$장의 카드에 쓰여진 정수의 최대 합을 구하라.

<접근 방법>

$N$개의 정수 중에서 $M$번의 "$B_j$장의 카드를 $C_j$로 변경"하는 연산을 수행한 뒤에 최대 합을 구하는 문제이다.

문제 예시를 살펴보자.

$$A = [5, 1, 4], \; B = [2, 1], \; C = [3, 5]$$

 

연산 과정을 차례대로 구성해보면 아래와 같이 수행될 것이다.

 

i) $B_1 = 2, \; C_1 = 3$

최대 2장 선택 가능하고, 그 중 A[1]을 선택하여 값을 변경하였다.

idx	0	1	2
A	5	1	4
연산 후 A	5	3	4

 

ii) $B_2 = 1, \; C_2 = 5$

최대 1장 선택 가능하고, 그 중 A[1]을 선택하여 값을 변경하였다.

idx	0	1	2
A	5	3	4
연산 후 A	5	5	4

 

따라서 $5 + 5 + 4 = 14$가 정답이 된다는 것을 알 수 있다.

 

해당 문제의 일반적인 해답을 찾아보자.

우선, $N$개의 정수 중 $C_j$보다 작은 숫자들을 대상으로 가장 작은 순서대로 최대 $B_j$개 변경하는 것을 생각할 수 있을 것이다. 왜냐하면 가장 작은 값들을 더 큰 값으로 변경해야 최댓값이 나올 것이기 때문이다.

 

하지만 일일이 하나씩 비교하는 것은 비효율적인 횟수가 많아 시간 초과가 발생할 수 있다. 아래의 경우를 고려해보자.

 

$$A = [1,2,3,4,5,6], \; B = [2,3], \; C = [3,4]$$

해당 상황에서 각 $j=1,2$에서 확인하는 숫자의 개수에 집중해보자.

$j=1$

idx	0	1	2	3	4	5
A	1	2	3	4	5	6
변경 대상 여부(O/X)	O	O	X	X	X	X
연산 후 A	3	3	3	4	5	6

 

$j=2$

idx	0	1	2	3	4	5
A	3	2	3	4	5	6
변경 대상 여부(O/X)	O	O	O	X	X	X
연산 후 A	4	4	4	4	5	6

 

즉, 단순 연산에서 확인하는 횟수는 총 $2 + 3 = 5$번이 된다. 해당 횟수를 더 줄일 수 있는 방법은 없을까?

 

해당 방법을 생각해보자. 어차피 모든 $B, C$에 대한 연산은 진행한다는 것은 확실하다. 그렇다면 B, C를 적용하는 순서는 변경되도 무방하다고 생각할 수 있다. 그렇다면 B, C를 B가 큰 순서대로 정렬한 다음, 연산을 진행하면 가장 작은 값들이 먼저 B의 가장 큰 값으로 변경되므로, 점점 B의 index를 진행할 수록 작은 수의 개수가 적어질 것이고, 비교 횟수가 더 작아질 것이다.

 

위의 예시에 한번 적용해보자.

$$A = [1,2,3,4,5,6], \; B = [2,3], \; C = [3,4] \rightarrow B = [3, 2], \; C = [4, 3] $$

 

$j=1$

idx	0	1	2	3	4	5
A	1	2	3	4	5	6
변경 대상 여부(O/X)	O	O	O	X	X	X
연산 후 A	4	4	4	4	5	6

 

$j=2$

idx	0	1	2	3	4	5
A	4	4	4	4	5	6
변경 대상 여부(O/X)	X	X	X	X	X	X
연산 후 A	4	4	4	4	5	6

 

즉, 확인 횟수는 총 $3 + 0 = 3$번이 되고, 실제로 횟수가 줄어든 것을 확인할 수 있다.

<소스 코드>

위의 접근법을 구현한 전체 소스 코드이다.

 

$N$개의 정수를 가지는 A 배열을 기반으로 PriorityQueue를 이용한 최소 heap을 구성하였다. $B, C$의 경우, tuple을 이용하여 change 배열에 삽입하였고, 해당 값을 B가 큰 tuple 순서로 정렬을 하고 연산을 진행하였다.

from queue import PriorityQueue

n, m = map(int, input().split())
*arr, = map(int, input().split())
hq = PriorityQueue()
[hq.put(i) for i in arr]
change = []
res = 0
for _ in range(m):
    change.append(tuple(map(int, input().split())))

change.sort(key=lambda x: -x[1])

for b, c in change:
    count = 1
    while count <= b:
        if not hq.empty():
            target = hq.get()
            if target > c:
                hq.put(target)
                break
            else:
                count += 1
                res += c
        else:
            break

while not hq.empty():
    res += hq.get()

print(res)