Pythonで木構造グラフにおける都市の最大人口を求めるプログラム

国をN個のノードとN-1本の辺からなる木構造として表現することを考えます。各ノードは町を表し、各辺は道路を表します。サイズN-1のリストsourceとdestが与えられ、i番目の道路はsource[i]とdest[i]を双方向に結んでいます。また、サイズNのリストpopulationも与えられ、population[i]はi番目の町の人口を表します。

ここで、いくつかの町を「都市」へアップグレードすることを考えます。ただし、以下の条件を満たす必要があります。

• 2つの都市が互いに隣接してはならない
• 町に隣接するすべてのノードは都市でなければならない(すべての道路は必ず町と都市をつなぐ)

これらの条件下で、すべての都市の人口合計の最大値を求めます。

例えば、入力が source = [2, 2, 1, 1]、dest = [1, 3, 4, 0]、population = [6, 8, 4, 3, 5] の場合、出力は15となります。ノード0、2、4を都市にアップグレードすると、人口の合計は 6 + 4 + 5 = 15 になるからです。

解法のアプローチ

この問題は、木構造を二部グラフとして2色に塗り分けることで解けます。木構造において、隣接するノード同士が必ず異なるグループに属するような塗り分け方は一意であるため、DFS(深さ優先探索)で各ノードを交互に2つのグループへ振り分け、それぞれのグループの人口合計を比較します。片方のグループを都市、もう片方を町とすれば条件は自動的に満たされるため、合計が大きい方を採用すればよいことになります。

具体的には、以下の手順で解きます。

• sourceとdestをもとに、グラフの隣接リストadjを作成する
• 関数dfs(x, choose)を定義する
• xがすでに訪問済みの場合は0を返す
• xを訪問済みとしてマークする
• ansを0で初期化する
• chooseがTrueの場合、ansにpopulation[x]を加算する
• adj[x]内の各近傍ノードneighborに対して、chooseを反転したdfs(neighbor, not choose)を呼び出し、その結果をansに加算する
• ansを返す
• メイン処理ではx = dfs(0, True)を計算し、max(x, sum(population) - x)を返す

以下の実装例を見ると、理解がより深まるでしょう。

実装例

from collections import defaultdict
class Solution:
    def solve(self, source, dest, population):
        adj = defaultdict(list)
        for a, b in zip(source, dest):
            adj[a].append(b)
            adj[b].append(a)

        seen = set()

        def dfs(x, choose):
            if x in seen:
                return 0
            seen.add(x)
            ans = 0
            if choose:
                ans += population[x]
            for neighbor in adj[x]:
                ans += dfs(neighbor, not choose)
            return ans

        x = dfs(0, True)
        return max(x, sum(population) - x)
        
ob = Solution()
source = [2, 2, 1, 1]
dest = [1, 3, 4, 0]
population = [6, 8, 4, 3, 5]
print(ob.solve(source, dest, population))

入力

[2, 2, 1, 1], [1, 3, 4, 0], [6, 8, 4, 3, 5]

出力

15