B - 商品選抜 / Product Selection 解説 by admin 印刷ページ
GPT 5.2 High概要
立ち入り可能なエリアにある商品のうち、販売価格が高いものから最大 \(K\) 個選んだときの合計金額の最大値を求めます。
考察
まず「立ち入り可能なエリアの商品のみが候補」という条件なので、各商品についてエリアが許可されているかを判定して候補を絞れます。問題は、候補商品の中から「最大 \(K\) 個」を選んで合計を最大化する方法です。
直感的には「候補を全部集めて降順ソートして上位 \(K\) 個を足す」で解けますが、候補が最大 \(N=2\times 10^5\) 個あるためソートは \(O(N\log N)\) で一応間に合うことも多いです。ただし、より一般に「上位 \(K\) 個だけ維持できれば十分」であり、毎回全体ソートする必要はありません。
ここで重要な気づきは次です: - 合計を最大にするには、候補の中から「価値が大きい順に \(K\) 個」を選ぶのが最適。 - つまり「上位 \(K\) 個」を効率よく管理できればよい。
そのために、サイズ \(K\) の集合を保ちつつ「いま選んでいる中で最小の値」をすぐ取り出せるデータ構造(最小ヒープ)を使います。
例:\(K=3\)、候補の価値が順に \([5,1,4,10,3]\) のとき
- 最初の3個で \(\{5,1,4\}\)(最小は1)
- 次の 10 は最小 1 より大きいので入れ替え → \(\{5,4,10\}\)
- 次の 3 は最小 4 以下なので無視
最終的に上位3個の合計が得られます。
アルゴリズム
- 立ち入り可能エリア集合 \(S\) を、配列
accessible[area](真偽値)として作る(\(S_j\) をTrueにする)。 - 最小ヒープ
heapと、現在ヒープに入っている値の合計totalを用意する。 - 全商品 \((V_i, C_i)\) について以下を行う:
accessible[C_i]がFalseならスキップ(候補外)。- まだ選んだ個数が \(K\) 未満なら、ヒープに \(V_i\) を追加し
totalに加算。 - すでに \(K\) 個選んでいるなら、ヒープ最小値
heap[0]と比較:- \(V_i > heap[0]\) なら、最小値と入れ替える(
heapreplace)。totalを差分だけ更新。 - そうでなければ、入れても上位 \(K\) 個に入らないので無視。
- \(V_i > heap[0]\) なら、最小値と入れ替える(
- 最後に
totalを出力する。
この方法で常に「これまで見た候補の中の上位 \(K\) 個」だけがヒープに残り、合計も同時に管理できます。
計算量
- 時間計算量: \(O(N \log K + M + Q)\)(実質的には各候補商品につきヒープ操作が最大1回で \(O(\log K)\))
- 空間計算量: \(O(M + K)\)(立ち入り可否配列とヒープ)
実装のポイント
立ち入り可否の判定は集合(
set)でもできますが、\(M\le 10^5\) なので真偽配列accessibleにすると高速です。合計値
totalを毎回更新しておくと、最後にヒープ内を合計し直す必要がありません。heapq.heapreplace(heap, v)は「最小要素を取り出してから push」より定数倍が小さく、差分更新もしやすいです。入力が大きいので
sys.stdin.buffer.read()でまとめて読むと高速です。ソースコード
import sys
import heapq
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
idx = 0
N = data[idx]; M = data[idx + 1]; K = data[idx + 2]; Q = data[idx + 3]
idx += 4
items_v = [0] * N
items_c = [0] * N
for i in range(N):
items_v[i] = data[idx]
items_c[i] = data[idx + 1]
idx += 2
accessible = [False] * (M + 1)
for _ in range(Q):
s = data[idx]
idx += 1
accessible[s] = True
heap = []
total = 0
for v, c in zip(items_v, items_c):
if not accessible[c]:
continue
if len(heap) < K:
heapq.heappush(heap, v)
total += v
else:
if v > heap[0]:
smallest = heapq.heapreplace(heap, v)
total += v - smallest
print(total)
if __name__ == "__main__":
main()
この解説は gpt-5.2-high によって生成されました。
投稿日時:
最終更新: