Pythonで全都市の市民が市場にアクセスできるようにする最小コストを求めるプログラム

問題の概要

n個の都市と、それらを結ぶm本の道路候補があるとします。市民が日用品を購入するには市場へのアクセスが必要ですが、現時点ではどの都市にも市場は存在せず、都市間の道路もまだ建設されていません。

2つの都市間に双方向の道路を建設できるのは、次の条件を満たす場合のみです。

• 片方の都市に市場が存在すること
• 市場のある地点から道路を経由してその都市へ到達できること

道路を1本建設するコストは x、市場を1つ建設するコストは y として与えられます。求めたいのは、すべての都市の市民が市場にアクセスできるようにするための最小コストです。配列 cities には、どの都市同士を道路で接続できるかという情報が含まれています。

入力例と出力例

たとえば、入力が n = 4、m = 3、x = 1、y = 2、cities = [[1, 2], [2, 3], [3, 4]] の場合、出力は 4 になります。

上の図のように、都市は1〜4の4つあります。都市1に市場を建設し、(1, 4) 間と (1, 3) 間に2本の道路を追加で建設すると、合計コストは 2 + 1 + 1 = 4 となり、これが最小コストになります。

解法の考え方

この問題は、グラフを連結成分単位で考えるとシンプルに解けます。

• x ≤ y の場合：道路よりも市場を建設する方が安いため、すべての都市に個別に市場を建設するのが最適です。総コストは n × x になります。
• x > y の場合：各連結成分ごとに市場を1つだけ建設し（コスト x）、残りの都市は道路で順番に接続していきます（1都市あたりコスト y）。連結成分の探索には幅優先探索（BFS）が便利です。

アルゴリズムの手順

1. x ≤ y ならば n × x を返します。
2. それ以外の場合は、隣接リスト adj_list を作成します。
3. 訪問済みフラグの配列 temp（初期値 True）と、答えとなる value = 0、両端キュー dq を用意します。
4. 各都市 cur について、未訪問なら市場を建設したものとみなして value に x を加算し、cur をキューに入れて BFS を開始します。
5. BFS で新しく到達した未訪問の都市ごとに、道路を建設したものとみなして value に y を加算します。
6. 最終的な value を返します。

各都市と各道路を一度ずつ処理するため、計算量は O(n + m) と非常に効率的です。

Pythonでの実装例

以下が実際の実装です。

from collections import defaultdict, deque

def solve(n, m, x, y, cities):
    if x <= y:
        # 市場の方が安い場合は全都市に市場を建設
        return n * x
    else:
        adj_list = defaultdict(list)
        for city in cities:
            city1 = city[0]
            city2 = city[1]
            adj_list[city1].append(city2)
            adj_list[city2].append(city1)

        temp = [True] * (n + 1)
        value = 0
        dq = deque()

        for cur in range(1, n + 1):
            if temp[cur]:
                # 新しい連結成分を見つけたら市場を建設
                value += x
                dq.append(cur)
                temp[cur] = False
                # BFSで同じ連結成分内の他の都市を道路で接続
                while dq:
                    for i in adj_list[dq.popleft()]:
                        if temp[i]:
                            dq.append(i)
                            temp[i] = False
                            value += y
        return value

print(solve(4, 3, 1, 2, [[1, 2], [2, 3], [3, 4]]))

入力

4, 3, 1, 2, [[1, 2], [2, 3], [3, 4]]

出力

4

まとめ

道路と市場のコストを比較し、x ≤ y なら全都市に市場を建設し、x > y なら連結成分ごとに市場を1つ建ててBFSで道路を接続していくことで、最小コストを効率的に求めることができます。