Pythonで家から最寄りのメールボックスまでの合計距離を最小化するプログラムの書き方

問題の概要

家の位置を表す配列 houses と整数 k が与えられます。houses[i] は一本の通り沿いにある i 番目の家の位置を示しており、この通りに k 個のメールボックスを設置することを考えます。このとき、各家から最寄りのメールボックスまでの距離の合計が最小になるような設置場所を求めるのが目的です。

たとえば、入力が houses = [6,7,9,16,22]、k = 2 の場合を考えてみましょう。メールボックスを位置 7 と 18 に設置すると、各家からの最短距離の合計は次のように計算できます。

|6−7| + |7−7| + |9−7| + |16−18| + |22−18| = 1 + 0 + 2 + 2 + 4 = 9

したがって、この場合の出力は 9 になります。

解決のアプローチ

この問題は、次の手順で解くことができます。

• リスト houses を昇順にソートします。
• 関数 util() を定義します。引数として区間の始点 idx、終点 n、使用できるメールボックスの数 k を受け取ります。
• k が 1 の場合（残りの家すべてにメールボックスを1つだけ割り当てる場合）：
  • 区間の中央値 core := houses[(n + idx) // 2] をメールボックスの設置位置とします。
  • 各区画の家との絶対差 |houses[i] - core| の総和を返します。
• それ以外の場合は result を無限大で初期化し、分割位置 i を idx から n まで動かしながら再帰的に探索します。
  • n - i < k - 1 となったら、残りのメールボックスを配置できなくなるためループを抜けます。
  • result を「左側の区間に1つのメールボックスを置いた場合」と「右側の区間に残り k-1 個を配置した場合」の和の最小値で更新します。
• 最終的に util(0, len(houses) - 1, k) を返します。

なぜ中央値を使うのか

1つのメールボックスで複数の家をカバーするとき、絶対誤差の合計を最小化する最適な設置位置は中央値です。外れ値の影響を受けにくく、平均値よりも合計距離を小さくできるため、このアルゴリズムでは各区間に対して中央値を採用しています。

実装例

以下のコードを実行して、実際の動作を確認してみましょう。

def solve(houses, k):
    houses.sort()
    def util(idx, n, k):
        if k == 1:
            core = houses[(n + idx) // 2]
            return sum([abs(houses[i] - core) for i in range(idx, n + 1)])
        result = float('inf')
        for i in range(idx, n + 1):
            if n - i < k - 1:
                break
            result = min(result, util(idx, i, 1) + util(i+1, n, k - 1))
        return result
    return util(0, len(houses) - 1, k)

houses = [6,7,9,16,22]
k = 2
print(solve(houses, k))

入力

[6,7,9,16,22], 2

出力

9

まとめ

この解法は、区間を再帰的に分割しながら、各区間の最適なメールボックス位置（中央値）を組み合わせて全体の最小コストを求める手法です。家の数やメールボックスの数が増えると計算量が増加するため、実運用ではメモ化などの動的計画法による高速化を検討するとよいでしょう。