Pythonでフードパケットを受け取れない人数を求めるプログラム

問題の概要

ある会議には、2種類の人々がいるとします。1つ目はベジタリアン(菜食)の昼食を希望する人々、もう1つは非菜食の昼食を希望する人々です。しかし、用意されているパケットの数は限られており、もし菜食希望者が非菜食のパケットを受け取った場合(またはその逆の場合)、その人はそのパケットを受け取らず、自分の希望するパケットが手に入るまで待ちます。

そこで、2種類のパケットと人々を、菜食を「0」、非菜食を「1」として表します。入力として、0と1で表されたn個のフードパケットを含む配列と、m人の行列(順番待ち)の希望を0と1で表した配列の2つが与えられます。もし誰かが自分の希望するパケットを受け取れなかった場合、その人は行列の最後尾に回り、再び自分の順番を待ちます。

私たちの目的は、フードパケットを受け取れない人数を求めることです。そうすることで、彼らの希望するパケットを追加で手配できるようになります。

入力例

例えば、people = [0,1,1,0]、packets = [0, 1, 0, 0] という入力の場合、出力は 1 になります。

この場合、非菜食を希望する人は2人いますが、非菜食のパケットは1つしかありません。行列の先頭にいる非菜食希望の人がそのパケットを受け取り、もう1人は他に非菜食のパケットがないため待ち続けます。したがって、出力は 1 となります。

解決の手順

この問題を解くために、以下の手順に従います。

• temp_arr を [0, 0] の新しいリストとして初期化します
• people 内の各 person について、temp_arr[person] を 1 増やします
• k を 0 で初期化します
• k が packets のサイズより小さい間、以下を繰り返します
  • temp_arr[packets[k]] が 0 より大きい場合、temp_arr[packets[k]] を 1 減らします
  • それ以外の場合、ループを抜けます
  • k を 1 増やします
• packets のサイズから k を引いた値を返します

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

def solve(people, packets):
    temp_arr = [0, 0]
    for person in people:
        temp_arr[person] += 1
    k = 0
    while k < len(packets):
        if temp_arr[packets[k]] > 0:
            temp_arr[packets[k]] -= 1
        else:
            break
        k += 1
    return len(packets) - k

print(solve([0, 1, 1, 0], [0, 1, 0, 0]))

入力

[0,1,1,0], [0, 1, 0, 0]

出力

1

アルゴリズムのポイント

この解法の鍵となるのは、行列の動きを実際にシミュレーションしない点です。各人が行列の最後尾に戻る動作を繰り返しても、結局のところ結果を左右するのは「各種類のパケットの数」と「各種類を希望する人の数」の対応関係だけだからです。

パケットを先頭から順番に配っていくとき、その時点で残っている該当種類の希望者数が0であれば、それ以降のパケットは誰にも受け取られません(残りの全員が別の種類を待ち続けるため)。したがって、ループが中断された位置以降のパケット数が、そのまま食べられない人数になります。

計算量は O(n + m)(n は人数、m はパケット数)であり、行列を一つずつ回す素朴なシミュレーションよりもはるかに効率的です。