【Python】ABC - 096 - Maximum Sum【heapq】

問題

3つの正整数$A, B, C$の中から一つ選び、2倍した値に書き換える。この操作を$K$回実施したときの合計の最大値を求める問題。

方針

1. 操作ごとに現在の3つの正整数から最大値を取り出し、2倍したものに書き換える
2. K回操作後の3つの正整数を合計する

工夫

蟻本　python プライオリティキュー　heapq　競技プログラミングのブログを参考に、heapqを用いて解く。

heapqを使うと良い場面として、

• 配列の要素が増える
• 最大値/最小値を取り出す

という操作を複数回繰り返すときに有効。

今回の問題では、プライオリティキューを用いてその時点での最大値を取り出し、その最大値を2倍したものを要素に追加する。
heapqの基本的な利用方法は、以下の通りである。Pythonのheapqを利用するときの注意点は、heapq.heappop()が最小値を取り出す仕様になっているので、最大値を取り出す操作を実施したい場合は、事前に要素にマイナスをかけたリストをプライオリティキューに変換する必要がある。

import heapq
heapq.heapify(a) # リストaのプライオリティキューに変換
heapq.heappush(a,x) # プライオリティキューに変換されたaに要素xを追加
heapq.heappop(a) # プライオリティキューに変換されたaから最小値を削除＆その最小値を出力

以下、実装を示す。

import heapq
A = list(map(int, input().split()))
A = [-a for a in A] # プライオリティキューで最大値を取り出す処理を実施するために正負反転
K = int(input())
heapq.heapify(A) # プライオリティキューに変換
for i in range(K):
    max_val = heapq.heappop(A) # 最大値の取得とプライオリティキューからの削除
    heapq.heappush(A, 2*max_val) # 取り出した最大値を2倍してプライオリティキューに追加
print(sum([-a for a in A]))