ABC141 D - Powerful Discount Tickets - ペンギンかっこいい

備忘録

問題

回答

import sys
import os
import math
import bisect
import collections
import itertools
import heapq

ii = lambda: int(sys.stdin.buffer.readline().rstrip())
il = lambda: list(map(int, sys.stdin.buffer.readline().split()))
fl = lambda: list(map(float, sys.stdin.buffer.readline().split()))
iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)]

iss = lambda: sys.stdin.buffer.readline().decode().rstrip()
sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split()))
isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)]

lcm = lambda x, y: x * y / math.gcd(x, y)

MOD = 10 ** 9 + 7
MAX = float('inf')


def main():
    if os.getenv("LOCAL"):
        sys.stdin = open("input.txt", "r")

    N, M = il()
    A = il()
    A = list(map(lambda x:x*-1, A))
    heapq.heapify(A)

    for _ in range(M):
        mn = heapq.heappop(A)*-1
        heapq.heappush(A, (mn//2)*-1)
    print(sum([a*-1 for a in A]))


if __name__ == '__main__':
    main()

考え方

割引券を使用するとき、
買う予定の品物の中から最も高い値段に対して使用する。
あとは上記の作業を割引券の枚数回（M回）行って、
最後に全ての品物の値段を合計することで回答を得ることができる。

実装に際して、
通常のlistで最も高い値をpop、割引券を使用（//2）した後、pushを行うと、
popとpushをM回行うことになる。
すると品物の数(A)は最大10 ** 9、割引券の枚数(M)は最大10 ** 5なので、
当然2000msには間に合わない。
そのため、優先度付きキューを使用する。

優先度付きキューの使い方は下記を参考にした。
【Python】優先度付きキューの使い方【heapq】 - Qiita

高速に最大値を得る方法さえわかれば、あとはシンプルな実装で回答を得ることが出来る。
注意する点は、heapqで最大値を得る際に、全ての値段を* -1しているので、
push時と回答時にそれを考慮すること。