【Python】2つの異なる都市へ同じ人数を送るときの最小コストを求めるプログラム

リスト costs があるとします。costs[i] は [c1, c2] という形式で表され、人 i が都市0へ移動するのに c1 のコストがかかり、都市1へ移動するのに c2 のコストがかかることを意味します。ここで、都市0と都市1へ同じ人数ずつ送りたい場合に必要な最小コストを求めるのがこの問題です。

たとえば、入力が costs = [[2, 6], [10, 3], [4, 9], [5, 8]] のとき、出力は 17 になります。これは、人0と人2を都市0へ、人1と人3を都市1へ割り当てると、都市0側のコストが 2 + 4 = 6、都市1側のコストが 3 + 8 = 11 となり、合計が 6 + 11 = 17 になるからです。

解法のアプローチ

この問題は貪欲法（グリーディ法）を使うことで効率的に解けます。まず全員を都市0へ送ると仮定してコストを合計し、その後「都市1へ切り替えた場合に追加でかかるコスト」（y − x）が最も小さい半分の人を都市1へ振り替えるという考え方です。

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

• s := 0 とする
• a := 空の新しいリストとする
• costs 内の各ペア (x, y) について以下を実行する
  • s := s + x
  • (y − x) をリスト a の末尾に追加する
• リスト a をソートする
• i を 0 から floor(a のサイズ ÷ 2) − 1 まで繰り返す
  • s := s + a[i]
• s を返す

実装例

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

def solve(costs):
    s = 0
    a = []
    for x, y in costs:
        s += x
        a += (y - x,)
    a.sort()
    for i in range(len(a) // 2):
        s += a[i]
    return s

costs = [[2, 6],[10, 3],[4, 9],[5, 8]]
print(solve(costs))

入力

[[2, 6],[10, 3],[4, 9],[5, 8]]

出力

17

計算量

時間計算量は O(n log n)（ソート処理が支配的）、空間計算量は O(n) です。n 人分の差分リストを作成してソートするだけで済むため、非常にシンプルかつ高速なアルゴリズムとなっています。