Submission #1501010

---

Source Code Expand

// {{{ Templates
#include <bits/stdc++.h>

#define show(x) cerr << #x << " = " << x << endl

using namespace std;
using ll = long long;
using pii = pair<int, int>;
using vi = vector<int>;

template <typename T>
ostream& operator<<(ostream& os, const vector<T>& v)
{
    os << "sz:" << v.size() << "\n[";
    for (const auto& p : v) {
        os << p << ",";
    }
    os << "]\n";
    return os;
}

template <typename S, typename T>
ostream& operator<<(ostream& os, const pair<S, T>& p)
{
    os << "(" << p.first << "," << p.second
       << ")";
    return os;
}


constexpr ll MOD = (ll)1e9 + 7LL;

template <typename T>
constexpr T INF = numeric_limits<T>::max() / 100;

// }}}

struct Graph {
    Graph(int n)
    {
        edge.resize(n);
    }
    void addEdge(int from, int to)
    {
        edge[from].push_back(to);
        edge[to].push_back(from);
    }
    vector<vector<int>> edge;
};

int num_dfs(const Graph& g, const int s, const vector<bool>& treasure, vector<int>& numarr, vector<bool>& visited)
{
    visited[s] = true;
    int num = 0;
    if (treasure[s]) {
        num++;
    }
    for (const int to : g.edge[s]) {
        if (visited[to]) {
            continue;
        }
        num += num_dfs(g, to, treasure, numarr, visited);
    }
    numarr[s] = num;
    return num;
}

int dfs(const Graph& g, const int s, const vector<int>& num, vector<bool>& visited)
{
    visited[s] = true;
    int length = 0;
    for (const int to : g.edge[s]) {
        if (visited[to]) {
            continue;
        }
        if (num[to] > 0) {
            length += dfs(g, to, num, visited) + 2;
        }
    }
    return length;
}

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    int n, x;
    cin >> n >> x;
    x--;
    Graph g(n);
    vector<bool> treasure(n, false);
    for (int i = 0; i < n; i++) {
        int a;
        cin >> a;
        treasure[i] = bool(a);
    }
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;
        a--, b--;
        g.addEdge(a, b);
    }
    vector<int> num(n, 0);
    vector<bool> visited(n, false);
    num_dfs(g, x, treasure, num, visited);
    fill(visited.begin(), visited.end(), false);
    cout << dfs(g, x, num, visited) << endl;

    return 0;
}

Submission Info

Judge Result