ARC030-B：ツリーグラフ

問題

https://beta.atcoder.jp/contests/arc030/tasks/arc030_2

サイズ n の木と数列 {h_n} が与えられる。 h_i = 1 となる頂点 i には宝石がある。頂点 x からはじめてすべての宝石を回収するのに必要な辺のコストを求めよ。

考え方

x を根として考える。宝石を回収するために頂点 x から dfs することを考える。 dfs で遷移する部分木に少なくとも 1 つ宝石が含まれていれば、その辺は必ず使う必要がある。辺は行きと帰りで 2 回使用するので、回収する際に必要な辺の数 * 2 で求めることができる。

       int[] h;
        List<Integer>[] g;
        int count;
        public void solve(int testNumber, InputReader in, PrintWriter out) {

            int n = in.nextInt(), x = in.nextInt()-1;
            h = in.nextIntArray(n);

            g = new ArrayList[n];
            g = Stream.generate(ArrayList::new).limit(n).toArray(List[]::new);

            for (int i = 0; i < n-1; i++) {
                int a = in.nextInt()-1, b = in.nextInt()-1;
                g[a].add(b);
                g[b].add(a);
            }

            dfs(x, -1);

            out.println(count * 2);

        }

        boolean dfs(int now, int par) {
            boolean ret = h[now] == 1;

            for (int to : g[now]) {
                if (to == par) continue;
                boolean child = dfs(to, now);
                if (child) count++;
                ret |= child;
            }

            return ret;
        }

ポイント

類題