予鈴

競プロとか備忘録とか…

ARC-B ツリーグラフ

B: ツリーグラフ - AtCoder Regular Contest 030 | AtCoder

n頂点からなるツリーグラフが与えられます。
いくつかの頂点にはお宝があり、そのお宝をすべて回収し、最初の頂点に戻ってくるまでの最小コストを求める問題です。

各頂点について考えます。
現在見ている頂点(vertex)を親に持つ頂点を通る必要がある→子孫のうち1つでも宝石を持っている。
という場合です。
宝石が見つかった場合は親にtrueを、そうでない場合はfalseを返すような再帰関数で解くことができます。

こういう問題が好きです()

#include <vector>
#include <iostream>
#include <utility>
#include <algorithm>
#include <string>
#include <deque>
#include <queue>
#include <tuple>
#include <queue>
#include <functional>
#include <cmath>
#include <iomanip>
#include <map>
#include <set>
#include <numeric>
#include <unordered_map>
#include <unordered_set>
#include <complex>
#include <iterator>
#include <array>
#include <memory>
#include <stack>
#define vi vector<int>
#define vvi vector<vector<int> >
#define ll long long int
#define vl vector<ll>
#define vvl vector<vector<ll>>
#define vb vector<bool>
#define vc vector<char>
#define vs vector<string>
#define ld long double
#define INF 1e9
#define EPS 0.0000000001
#define rep(i,n) for(int i=0;i<n;i++)
#define loop(i,s,n) for(int i=s;i<n;i++)
#define all(in) in.begin(), in.end()
template<class T, class S> void cmin(T &a, const S &b) { if (a > b)a = b; }
template<class T, class S> void cmax(T &a, const S &b) { if (a < b)a = b; }
#define MAX 9999999
using namespace std;
class Edge{
public:
    int to,from,cost;
};
vector<int>jem(101,0);
vector<vector<int>>edge(101);
vector<int>cost(101,0);
bool  dfs(int from,int vertex){
    bool var=false;
    if(edge[vertex].size()==1&&jem[vertex]==0)return false;
    rep(i,edge[vertex].size()){
        int next=edge[vertex][i];
        if(next==from)continue;
        if(dfs(vertex,next))var=true;
    }
    if(var||jem[vertex])cost[from]+=cost[vertex]+2;
    else return false;
    return true;
}
 
int main(){
    int  n,x; cin>>n>>x;
    rep(i,n){
        int foo; cin>>foo;
        if(foo)jem[i+1]=1;
    }
    rep(i,n-1){
        int to,from; cin>>to>>from;
        edge[to].push_back(from);
        edge[from].push_back(to);
    }
    rep(i,edge[x].size())dfs(x,edge[x][i]);
    cout<<cost[x]<<endl;
 
}