Pythonで幸せにできる顧客の最大人数を求めるプログラム

問題概要

同じ長さを持つ2つのリスト customers(毎分の来店者数)と mood(客の機嫌)、および整数 k が与えられます。毎分 i には customers[i] 人の客が店を訪れ、mood[i] = 1 のときはその客たちが幸せであること、mood[i] = 0 のときは悲しいことを表します。ここで、連続する長さ k の区間をひとつ選び、その間の mood をすべて 1 に変更できます。最終的に幸せにできる人数の最大値を求めましょう。

例として、入力が customers = [2, 3, 6, 6, 3]、mood = [1, 1, 0, 0, 0]、k = 2 の場合を考えてみます。mood[2] と mood[3] を 1 に変更すれば、2 + 3 + 6 + 6 = 17 人の客を幸せにできるため、答えは 17 になります。

解法の考え方

この問題は、累積和(プレフィックスサム)とスライディングウィンドウを組み合わせることで、O(n) という高速な計算量で解くことができます。手順は以下のとおりです。

• n を mood の長さとします。
• 長さ n + 1 の累積和配列 a を 0 で初期化します。
• 変数 s を 0 にします。
• i を 0 から n - 1 まで順に処理します。
  • a[i + 1] = a[i] とします。
  • mood[i] が 1(非ゼロ)の場合:s に customers[i] を加算します。
  • それ以外の場合(mood[i] が 0):累積和配列側に a[i + 1] へ customers[i] を加算します。
• 変数 d を 0 にします。
• i を k から n まで動かしながら、d を d と a[i] - a[i - k] の最大値で更新します。これは、長さ k の区間内に含まれる「悲しい客」の合計の最大値を求める操作です。
• s + d を返します。

つまり、「もともと幸せな客の総数 s」に、「長さ k の窓で救済できる悲しい客の最大人数 d」を足したものが答えになります。

実装例

以下は Python による実装例です。

def solve(customers, mood, k):
    n = len(mood)
    a = [0] * (n + 1)
    s = 0
    for i in range(n):
        a[i + 1] = a[i]
        if mood[i]:
            s += customers[i]
        else:
            a[i + 1] += customers[i]
    d = 0
    for i in range(k, n + 1):
        d = max(d, a[i] - a[i - k])
    return s + d

customers = [2, 3, 6, 6, 3]
mood = [1, 1, 0, 0, 0]
k = 2
print(solve(customers, mood, k))

入力

[2, 3, 6, 6, 3], [1, 1, 0, 0, 0], 2

出力

17

計算量

どちらのループもデータを一度ずつ走査するだけなので、時間計算量は O(n)、追加で必要なメモリは累積和配列ぶんの O(n) です。全探索のように O(n × k) かかる実装と比べても非常に効率的であり、入力サイズが大きくなっても安定して動作します。