C++で木の中のすべてのリンゴを収集するための最小時間を求める

問題概要

n個の頂点からなる無向木を考えます。頂点には0からn-1までの番号が付けられており、いくつかの頂点にはリンゴが置かれています。木の1つの辺を移動するのに1秒かかるとき、頂点0から出発してすべてのリンゴを集め、再び頂点0に戻るまでに必要な最小時間（秒）を求めてください。

無向木の辺は配列 edges として与えられ、edges[i] = [from_i, to_i] は頂点 from_i と頂点 to_i を結ぶ辺が存在することを表します。さらに、hasApple というブール値の配列も与えられ、hasApple[i] = true の場合は頂点 i にリンゴが存在し、false の場合は存在しないことを意味します。

たとえば、入力が n = 7、edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]]、hasApple = [false, false, true, false, true, true, false] のとき、出力は 8 になります。

上図のように、赤い頂点はリンゴが存在する頂点を示しています。緑の矢印は、すべてのリンゴを収集するための最適な経路の一例です。

解法のアプローチ

この問題の鍵となるのは、「リンゴが存在する部分木につながる辺は、必ず行きと帰りの2回通過する必要がある」という点です。したがって、DFS（深さ優先探索）を使って各部分木にリンゴが含まれているかどうかを判定し、含まれる場合にはその辺のコストを2倍して合計すれば、答えを求められます。具体的な手順は以下の通りです。

• 関数 dfs(node, par, a, graph) を定義します。node は現在の頂点、par は親の頂点、a はリンゴの有無を格納した配列、graph は隣接リストです。
• temp := 0 で初期化します。
• graph[node] 内の各要素 x に対して、次の処理を繰り返します。
  • x が par と等しい場合は、後続の処理をスキップして次の反復へ進みます。
  • それ以外の場合は、temp := temp + dfs(x, node, a, graph) とします。
• ret := ret + temp * 2 として、結果に加算します。
• a[node] + temp > 0 であれば true を、そうでなければ false を返します。
• メイン処理では、以下を実行します。
  • ret := 0 で初期化します。
  • n 個のリストからなる隣接リスト graph を用意します。
  • i := 0 から e のサイズ未満の間、i を1ずつ増やしながら、次を繰り返します。
    • graph[e[i][0]] の末尾に e[i][1] を追加します。
    • graph[e[i][1]] の末尾に e[i][0] を追加します。
  • dfs(0, -1, a, graph) を呼び出します。
  • ret を返します。

C++実装例

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

#include <bits/stdc++.h>
using namespace std;
const int N = 1e6;
class Solution {
public:
    set<int> visited;
    int ret;
    int dfs(int node, int par, vector<bool>& a, vector<int> graph[]){
        int temp = 0;
        for (int x : graph[node]) {
            if (x == par)
                continue;
            temp += dfs(x, node, a, graph);
        }
        ret += temp * 2;
        return a[node] + temp > 0;
    }
    int minTime(int n, vector<vector<int> >& e, vector<bool>& a){
        ret = 0;
        vector<int> graph[n];
        for (int i = 0; i < e.size(); i++) {
            graph[e[i][0]].push_back(e[i][1]);
            graph[e[i][1]].push_back(e[i][0]);
        }
        dfs(0, -1, a, graph);
        return ret;
    }
};
main(){
    Solution ob;
    vector<vector<int>> v = {{0,1},{0,2},{1,4},{1,5},{2,3},{2,6}};
    vector<bool> v1 = {false,false,true,false,true,true,false};
    cout << (ob.minTime(7,v, v1));
}

入力

7, {{0,1},{0,2},{1,4},{1,5},{2,3},{2,6}},
{false,false,true,false,true,true,false}

出力

8

計算量

このアルゴリズムの時間計算量は O(n)、空間計算量は O(n) です。各頂点を一度だけ訪問するため、大きな木構造に対しても効率的に動作します。