【Python】腐る前に食べられるリンゴの最大数を求めるアルゴリズム(最小ヒープで解く)

同じ長さ n の2つの配列 days と apples があるとします。ある特別なリンゴの木が n 日間連続でリンゴを実らせます。i 日目には apples[i] 個のリンゴがなり、それらは days[i] 日後に腐ります。言い換えると、i + days[i] 日目にはそのリンゴは腐って食べられなくなります。また、apples[i] = 0 かつ days[i] = 0 の場合は、i 日目には新しいリンゴが実らないことを表します。

1日に食べられるリンゴは最大1個です(最初の n 日が過ぎても、腐っていないリンゴが残っていれば食べ続けられます)。このとき、最終的に食べられるリンゴの最大数を求めるのがこの問題の目的です。

例で理解する

入力が apples = [1,2,3,5,2]、days = [3,2,1,4,2] の場合、出力は 7 になります。食べ方の内訳は次のとおりです。

• 1日目: 1日目に実ったリンゴを1個食べます。
• 2日目: 2日目に実ったリンゴを1個食べます。
• 3日目: 2日目に実った残りのリンゴを1個食べます。この日を過ぎると、3日目に実ったリンゴは腐ってしまいます。
• 4日目〜7日目: 4日目に実ったリンゴを毎日1個ずつ、合計4個食べます。

解法のアプローチ

この問題は貪欲法と最小ヒープ(優先度付きキュー)を組み合わせることで効率的に解けます。鍵となる発想は、「腐る期限が近いリンゴから優先的に食べる」ことです。ヒープには「腐敗日」と「残り個数」のペアを格納し、常に期限が最も早いリンゴを先に消費していきます。

具体的な手順は以下のとおりです。

1. 空の最小ヒープ minheap を用意します。
2. day := 0、res := 0 で初期化します。
3. i を 0 から apples のサイズ - 1 まで繰り返します。
   • day := i とします。
   • ヒープが空でなく、先頭要素の腐敗日が day 未満である間、先頭要素を削除します(すでに腐ったリンゴの掃除)。
   • nbrApple := apples[i]、expiration := i + days[i] - 1 を計算します。
   • nbrApple > 0 の場合、ペア (expiration, nbrApple) をヒープに挿入します。
   • ヒープが空でなければ、先頭の (date, apple) を取り出して res を1増やします。apple > 1 なら (date, apple-1) をヒープに戻します。
4. n 日以降も、ヒープが空になるまで1日ずつ日を進めながら、腐っていないリンゴを同じ要領で食べ続けます。
5. 最後に res を返します。

実装例(Python)

それでは、実際のコードを見てみましょう。

import heapq
def solve(apples, days):
    minheap = []
    heapq.heapify(minheap)
    day = 0
    res = 0
    for i in range(len(apples)):
        day = i

        while minheap and minheap[0][0] < day:
            heapq.heappop(minheap)

        nbrApple = apples[i]
        expiration = i + days[i]-1

        if nbrApple > 0:
            heapq.heappush(minheap, (expiration, nbrApple))

        if minheap:
            date, apple = heapq.heappop(minheap)
            res += 1
            if apple > 1:
                heapq.heappush(minheap, (date, apple-1))

    while minheap:
        day += 1
        while minheap and minheap[0][0] < day:
            heapq.heappop(minheap)
        if minheap == []:
            break
        date, apple = heapq.heappop(minheap)
        res += 1
        if apple > 1:
            heapq.heappush(minheap, (date, apple-1))

    return res

apples = [1,2,3,5,2]
days = [3,2,1,4,2]
print(solve(apples, days))

入力

[1,2,3,5,2],[3,2,1,4,2]

出力

7

計算量

各リンゴはヒープへの挿入と削除を高々1回ずつしか行わないため、時間計算量は O((N + M) log N)(N は日数、M は食べられるリンゴの総数)、空間計算量は O(N) となります。貪欲に「期限の近いリンゴから食べる」ことで、無駄なく最大数を達成できる点がこのアルゴリズムのポイントです。