AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC160 E - Red and Green Apples

問題原文

atcoder.jp

問題要旨

 A 個の赤いリンゴ、  B 個の緑のリンゴ、  C 個の無色のリンゴがある。
それぞれの美味しさは、  p_i,  q_j,  r_k である。
また、 X 個の赤いリンゴ、  Y 個の緑のリンゴを食べることを考える。
食べるリンゴの美味しさの合計を最大化せよ。ただし、無色のリンゴは赤いリンゴとしても緑のリンゴとしても使える。

  •   1 \leq X \leq A \leq 10^{5}
  •   1 \leq Y \leq B \leq 10^{5}
  •   1 \leq C \leq 10^{5}

解法

こういうのは逆の手順を考えると途端に見通しがよくなる。
つまり、「あらかじめ各無色のリンゴを赤/緑として食べることにする」のではなく、「赤/緑としてたべたことにする」、という視点で考える。

すると、「先に赤いリンゴの美味しさが大きい方から  X 個、 緑のリンゴの美味しさが大きい方から  Y 個 とっておいて、あとから価値が低いものを無色のリンゴと置き換えてその色にしたことにする」、という操作が浮かぶ。

価値が最少のものを取り出すのには、優先度付きキューを使う。

実装

import heapq
X, Y, A, B, C = map(int, input().split())
P = sorted(list(map(int, input().split())), reverse=True)
Q = sorted(list(map(int, input().split())), reverse=True)
R = sorted(list(map(int, input().split())), reverse=True)

hq = P[:X] + Q[:Y]
heapq.heapify(hq)

for r in R:
    heapq.heappushpop(hq, r)

print(sum(hq))

感想

問題文の言い換えで見通しがよくなるシリーズすきです。