Submission #1455026

---

Source Code Expand

#include <bits/stdc++.h>
#define rep(n) for( int I = 0; (I) < (n); ++(I) )
#define repeat(i, n) for( int i = 0; (i) < (n); ++(i) )
#define repeat_to(i, n) for( int i = 0; (i) <= (n); ++(i) )
#define repeat_from(i, m, n) for( int i = (m); (i) < (n); ++(i) )
#define repeat_from_to(i, m, n) for( int i = (m); (i) <= (n); ++(i) )
#define repeat_from_reverse(i, m, n) for( int i = (n) - 1; (i) >= (m); --(i) )
#define dump(x) cout << " " << #x << "=" << x
#define vdump(v) for(size_t t=0; t<v.size(); ++t){cout << " " << #v << "[" << t << "]=" << v[t];} cout << endl
using namespace std;
using lint = long long;
using ld = long double;

/****************** グラフ ******************/
// 諸宣言
using Weight = lint;
using Flow = lint;
struct Edge {
    int src, dst;
    Weight weight;
    Flow cap;
    Edge() : src(0), dst(0), weight(0) {}
    Edge(int s, int d, Weight w) : src(s), dst(d), weight(w) {}
};
using Edges = std::vector<Edge>;
using Graph = std::vector<Edges>;
using Array = std::vector<Weight>;
using Matrix = std::vector<Array>;

void add_edge(Graph &g, int a, int b, Weight w = 1) {
    g[a].emplace_back(a, b, w);
    g[b].emplace_back(b, a, w);
}
void add_arc(Graph &g, int a, int b, Weight w = 1) { g[a].emplace_back(a, b, w); }

int main(void) {
    int n, x; cin >> n >> x; --x;
    vector<bool> treasure(n);
    repeat(i, n) {
        int t; cin >> t;
        if(t == 0) treasure[i] = false;
        else treasure[i] = true;
    }

    Graph g(n);
    rep(n-1) {
        int a, b; cin >> a >> b;
        --a; --b;
        add_arc(g, a, b);
    }

    function<pair<bool, int>(Graph&, int)> rec = [&treasure, &rec](Graph &g, int node) -> pair<bool, int> {
        
        //cout << "g[" << node << "].size()=" << g[node].size() << endl;
        
        if (g[node].size() == 0) {
            if (treasure[node]) return {true, 0};
            else return {false, 0};
        }
        
        vector<pair<bool,int>> res;
        repeat(i, (int)g[node].size()) {
            res.push_back( rec(g, g[node][i].dst) );
        }

        int ans = 0;
        for (auto e: res) {
            if (e.first) ans += (e.second + 2);
        }
        if (ans == 0) {
            return {false, 0};
        }
        else {
            return {true, ans};
        }
    };

    cout << rec(g, x).second << endl; 
    return 0;

}

Submission Info

Judge Result