Pythonで色付きボールの販売による最大利益を求めるプログラムの作成方法

問題の概要

ここに inventory という配列があるとします。inventory[i] は、i 番目の色のボールの初期在庫数を表します。さらに、顧客が購入したいボールの総数を表す値 orders も与えられます。ボールはどの順序でも販売でき、顧客はどんな色のボールでも受け入れます。

このボールの価値には特別なルールがあります。各色のボールの価値は、「その色のボールが現在インベントリに何個残っているか」に等しくなります。たとえば、現在青いボールが6個ある場合、最初の1個は価格6で売れ、残りが5個になるため、次の青いボールは価格5で売れることになります。この条件のもとで、orders 個のボールを売った後に得られる合計価値の最大値を求めます。答えが非常に大きくなる可能性があるため、10^9 + 7 で割った余りを返してください。

入力例

たとえば、入力が inventory = [5,7]、orders = 6 の場合、出力は 31 になります。これは、1つ目の色のボールを2回(価格 5, 4)、2つ目の色のボールを4回(価格 7, 6, 5, 4)売ることができ、合計利益が 5+4+7+6+5+4 = 31 となるからです。

解法のアプローチ

この問題を効率的に解くには、二分探索(バイナリサーチ)を活用します。高い価格のボールから順に売っていくのが最適戦略であるため、「これ以下の価格では売らない」というしきい値を二分探索で特定します。手順は以下の通りです。

• low := 0、high := 10000000 と初期化します。

• low < high の間、次を繰り返します。

  • mid := (low + high) // 2(商)とします。

  • s := 0 とします。

  • inventory 内の各要素 i について、i > mid ならば s := s + (i - mid) とします。

  • s > orders ならば low := mid + 1、そうでなければ high := mid とします。

• 再度 mid := (low + high) // 2 を計算します。

• ans := 0 とします。

• inventory 内の各要素 i について、i > mid ならば次を行います。

  • ans := ans + i*(i+1)//2 - mid*(mid+1)//2(連続する整数の総和の公式を利用)

  • orders := orders - (i - mid)

• ans := ans + orders * mid を加算します。

• ans mod (10^9 + 7) を返します。

Pythonでの実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(inventory, orders):
    low = 0
    high = 10000000

    # 売却しきい値を二分探索で求める
    while low < high:
        mid = (low + high) // 2

        s = 0
        for i in inventory:
            if i > mid:
                s += i - mid

        if s > orders:
            low = mid + 1
        else:
            high = mid

    mid = (low + high) // 2

    ans = 0
    for i in inventory:
        if i > mid:
            ans += i * (i + 1) // 2 - mid * (mid + 1) // 2
            orders -= i - mid

    ans += orders * mid
    return ans % (10**9 + 7)

inventory = [5, 7]
orders = 6
print(solve(inventory, orders))

入力

[5, 7], 6

出力

31

アルゴリズムのポイント

この解法の鍵となるのは、次の2点です。

• 二分探索によるしきい値の特定: 「mid 以上の在庫を持つ色から、mid を超える分だけ売る」という戦略で販売可能な個数を数え、orders と比較しながらしきい値を絞り込みます。

• 等差数列の和の公式: 在庫 i 個から mid 個まで売るときの売上は i + (i-1) + … + (mid+1) となり、これは i*(i+1)/2 − mid*(mid+1)/2 として定数時間で計算できます。

計算量は、二分探索の各反復で在庫配列を走査するため、O(n log m)(n は色の数、m は在庫の最大値)となります。貪欲法的な発想と二分探索を組み合わせることで、大量の注文にも高速に対応できるのが特徴です。