Submission #1113790

---

Source Code Expand

#include <bits/stdc++.h>
using namespace std;
const int size = 100;

int N, X;
int H[size];
vector<int> edges[size];
bool hasJewel[size];
int cnt[size];
vector<int> idByDepth;

void init() {
  cin >> N >> X;
  X -= 1;
  for(int i = 0; i < N; i++) {
    int a;
    cin >> a;
    if(a == 1)
      hasJewel[i] = true;
    else
      hasJewel[i] = false;
  }
  
  for(int i = 0; i < N - 1; i++) {
    int a, b;
    cin >> a >> b;
    a -= 1;
    b -= 1;
    edges[a].push_back(b);
    edges[b].push_back(a);
  }
}

int parent[size];

void getparent() {
  bool visited[N];
  for(int i = 0; i < N; i++)
    visited[i] = false;
  queue<int> q;
  q.push(X);
  parent[X] = -1;
  visited[X] = true;
  while(q.empty() == false) {
    int from = q.front();
    q.pop();
    for(int to: edges[from]) {
      if(visited[to] == false) {
	idByDepth.push_back(to);
	visited[to] = true;
	parent[to] = from;
	q.push(to);
      }
    }
  }
  reverse(idByDepth.begin(), idByDepth.end());

}

void f() {


  for(int i: idByDepth) {
    if(hasJewel[i]) {
      cnt[ parent[i] ] += (cnt[i] + 1);
      hasJewel[ parent[i] ] = true;
      hasJewel[i] = false;
      f();
    }
  }
}

void solve() {
  getparent();
  f();
  cout << cnt[X] * 2 << endl;
}

int main() {
  init();
  solve();
}

Submission Info

Judge Result