Official

D - 果樹園の収穫 / Orchard Harvest Editorial by admin

Qwen3-Coder-480B

概要

\(N\) 種類の果樹から最大 \(M\) 回収穫を行い、得られる果物の総数を最大化する問題。

考察

各果樹から収穫するたびに得られる果物の数が減っていく(\(\max(F_i - (k-1) \times D_i, 0)\))。
そのため、単純に「初回収穫量 \(F_i\) が大きい順に収穫する」だけでは最適にならない。

例えば、以下のようなケースを考える:

  • 木 A: \(F_A = 10\), \(D_A = 9\)
  • 木 B: \(F_B = 8\), \(D_B = 1\)
  • 収穫回数の上限 \(M = 2\)

木 A を2回収穫すると、\(10 + (10 - 9) = 11\) 個を得られる。
一方、木 B を2回収穫すると、\(8 + (8 - 1) = 15\) 個を得られる。
このように、疲弊度が小さい木は長期的に見てお得である可能性がある。

したがって、毎回「次に最も多くの果物を得られる収穫」を選ぶのが最適である。これは貪欲法で実現可能。

TLE/WAを避けるために

全パターンを試すと組み合わせが爆発するので現実的ではない。
また、単純に木を何回も舐めても、次にどの木から収穫すべきかが分からない。
そこで、「次に最も多く得られる果物の個数」を高速に取得・更新するために優先度付きキュー(ヒープ)を使う。

アルゴリズム

  1. 各果樹について、初回収穫量をヒープに登録する(ただし、収穫量が0より大きい場合のみ)。
  2. ヒープは「次に最も多く得られる果物の個数」が常に取り出せるように、最大ヒープとして使用する(Pythonでは最小ヒープしかないため、符号を反転して格納)。
  3. \(M\) 回の収穫のうち、毎回以下を行う:
    • ヒープから最も大きな値を取り出し、その分を合計に加える。
    • その木の次の収穫量を計算し、それが正であれば再びヒープに追加する。
  4. ヒープが空になったらそれ以上収穫できないので終了。

これにより、常に「今最も得られる果物が多い収穫」を選んでいく貪欲法となる。

計算量

  • 時間計算量: \(O((N + M) \log N)\)
    • 最初に \(N\) 個の要素をヒープに挿入 → \(O(N \log N)\)
    • \(M\) 回の操作で、毎回最大1つ要素を削除・挿入 → \(O(M \log N)\)
  • 空間計算量: \(O(N)\)
    • ヒープに最大 \(N\) 個の要素を持つ

実装のポイント

  • Python の heapq は最小ヒープなので、最大値を取り出すためには要素の符号を反転させる必要がある。
  • 収穫量が 0 になったら、それ以上その木から収穫しても意味がないため、再度ヒープに追加しない。
  • 入力を高速に読み込むために sys.stdin.read を使用している。

ソースコード

import heapq

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    N = int(data[0])
    M = int(data[1])
    
    fruits = []
    index = 2
    for _ in range(N):
        F = int(data[index])
        D = int(data[index+1])
        index += 2
        fruits.append((F, D))
    
    # 各木について、次の収穫で得られる果物の数を管理するための優先度付きキュー
    # 最大値を取り出すために符号を反転させて格納
    heap = []
    
    # 初期状態で各木の1回目の収穫をキューに追加
    for F, D in fruits:
        if F > 0:
            heapq.heappush(heap, (-F, F, D))  # (-現在の収穫量, 初回収穫量, 疲弊度)
    
    total_fruits = 0
    
    # M回まで収穫を行う
    for _ in range(M):
        if not heap:
            break
        
        # 現在最も多くの果物が得られる収穫を選択
        neg_current, current, D = heapq.heappop(heap)
        total_fruits += current
        
        # 次の収穫での果物の数を計算
        next_amount = max(current - D, 0)
        
        # 次の収穫が有効であればキューに追加
        if next_amount > 0:
            heapq.heappush(heap, (-next_amount, next_amount, D))
    
    print(total_fruits)

if __name__ == "__main__":
    main()

この解説は qwen3-coder-480b によって生成されました。

posted:
last update: