Submission #1124077

---

Source Code Expand

#include <iostream>
#include <vector>
#include <set>
#include <algorithm>
#include <iomanip>
using namespace std;
using ll = long long;

constexpr ll M = 1e9+7;

int dfs(vector<vector<int>>& g, int v, vector<bool>& used, vector<bool>& dia) {
    used[v] = true;
    int res = 0;
    for(int i=0; i<g[v].size(); ++i) {
        if(!used[g[v][i]]) {
            int t = dfs(g, g[v][i], used, dia);
            if(t > 0) {
                res += 2 + (t == 1 ? 0 : t);
            }
        }
    }
    return (res == 0 && dia[v] ? 1 : res);
}

int main() {
    int n, x;
    cin >> n >> x;
    x--;
    vector<vector<int>> g(n);
    vector<bool> dia(n, false);
    for(int i=0; i<n; ++i) {
        int h; cin >> h;
        dia[i] = (h == 1);
    }
    dia[x] = false;
    for(int i=0; i<n-1; ++i) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    vector<bool> used(n, false);
    cout << dfs(g, x, used, dia) << endl;
}

Submission Info

Judge Result