Submission #19918292

---

Source Code Expand

#include <atcoder/all>
using namespace atcoder;
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const double pi = 3.14159265359;
const ll INF = 1LL << 60;
template<class T> inline bool chmin(T& a, T b){
  if (a > b){
    a = b;
    return true;
  }
  return false;
}
template<class T> inline bool chmax(T& a, T b){
  if (a < b){
    a = b;
    return true;
  }
  return false;
}

vector<int> hasStone;
vector<vector<int>> edges;

int dfs(int par, int i){
  auto& es = edges[i];

  int dist = 0;
  int found = false;
  for (int j : es){
    if (j != par){
      int r = dfs(i, j);
      if (r >= 0){
	found = true;
	dist += r + 2;
      }
    }
  }

  if (found){
    return dist;
  } else if (hasStone[i]) {
    return 0;
  } else {
    return -1;
  }
}

int main()
{
  int n, x;
  cin >> n >> x; x--;

  hasStone.resize(n);
  for (int& b : hasStone) cin >> b;

  edges.resize(n);
  for (int i = 0; i < n-1; i++){
    int a, b;
    cin >> a >> b; a--; b--;
    edges[a].push_back(b);
    edges[b].push_back(a);
  }

  int ans = dfs(-1, x);
  cout << max(0, ans) << endl;
  return 0;
}

Submission Info

Judge Result