ARC030 B - ツリーグラフ

B - ツリーグラフ

概要

n頂点の木が与えられる。いくつかの頂点には宝石が置いてある。
頂点xからスタートして辺をたどり頂点にある宝石をすべて回収してxに戻ってきたい。
移動経路の最短の長さを答えよ。

制約

1 ≦ n ≦ 100

方針

xを根として考え、xから宝石のあるノードに順番に移動してゆく。
共通の親を持つ子については、その親(lca)を拠点にして移動するイメージ。

下記のDFSを行う。

• 葉なら0を返す
• 頂点qのそれぞれの子uについて下記を合計する
  • uが数値yを返せば、y+2を加算（q→uと移動、uからuの子に移動（移動距離はy）、u→qと移動)
  • uに宝石があれば、2を加算（q→u→qと移動）

解答

Submission #4705535 - AtCoder Regular Contest 030

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
 
const int max_n = 103;
int n, x;
vector<vector<int>> g(max_n, vector<int>());
 
int h[max_n];
 
int dfs(int p, int q){
	// cout << "in: " << p << "," << q << endl;
 
	int ret = 0;
	for(int u : g[q]){
		if(u == p) continue;
 
		int y = dfs(q, u);
		// cout << q << "," << u << "," << h[u] << "," << y << endl;
 
		if(y > 0){
			ret += y + 2;
		}else if(h[u] == 1){
			ret += 2;
		}
	}
 
	// cout << "out: " << p << "," << q << "," << ret << endl;
	return ret;
}
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	cin >> n >> x;
 
	rep(i, n) cin >> h[i];
 
	rep(i, n-1){
		int a, b;
		cin >> a >> b;
		g[a-1].push_back(b-1);
		g[b-1].push_back(a-1);
	}
 
	cout << dfs(-1, x-1) << endl;
}