Submission #1476621

---

Source Code Expand

#include <iostream>
#include <queue>
#include <vector>
#include <cmath>
#include <algorithm>
#include <map>

using namespace std;
using ll = long long;
const ll INF = 1e9;

int N;
vector<bool> H(101, false);
map<int, vector<int>> G;
map<int, vector<int>> U;

bool decide(int cur, int parent) {
    bool use_cur = false;
    for (auto i : G[cur]) {
        if (i == parent) continue;
        if (H[i]) {
            U[cur].push_back(i);
            use_cur = true;
        }
        if (decide(i, cur)) {
            use_cur = true;
        }
    }
    if (use_cur) {
        U[parent].push_back(cur);
    }
    return use_cur;
}

int steps(int cur, int parent) {
    int ret = 0;
    for (auto i : G[cur]) {
        if (i == parent) continue;
        for (auto j : U[cur]) {
            if (i == j) {
                ret += 2;
                break;
            }
        }
        ret += steps(i, cur);
    }
    return ret;
}

int main() {
    int X;
    cin >> N >> X;
    H.resize(N+1);
    for (int i = 0; i < N; i++) {
        int a;
        cin >> a;
        H[i+1] = (a == 1);
    }
    for (int i = 0; i < N - 1; i++) {
        int a, b;
        cin >> a >> b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    decide(X, -1);
    cout << steps(X, -1) << endl;

    return 0;
}

Submission Info

Judge Result