Submission #1321944

---

Source Code Expand

#include <bits/stdc++.h>

#ifndef LOCAL_
#define fprintf if( false ) fprintf
#endif // LOCAL_
// #define dump() fprintf(stderr, "#%s.%d\n", __func__, __LINE__);
#define dumpl(x1) fprintf(stderr, "#%s.%d (%s) = (%ld)\n", __func__, __LINE__, #x1, x1);
#define dumpll(x1, x2) fprintf(stderr, "#%s.%d (%s, %s) = (%ld, %ld)\n", __func__, __LINE__, #x1, #x2, x1, x2);
#define dumplll(x1, x2, x3) fprintf(stderr, "#%s.%d (%s, %s, %s) = (%ld, %ld, %ld)\n", __func__, __LINE__, #x1, #x2, #x3, x1, x2, x3);
#define dumpd(x1) fprintf(stderr, "#%s.%d (%s) = (%lf)\n", __func__, __LINE__, #x1, x1);
#define dumpdd(x1, x2) fprintf(stderr, "#%s.%d (%s, %s) = (%lf, %lf)\n", __func__, __LINE__, #x1, #x2, x1, x2);
#define loop for(;;)

template<typename T> void scan1(T& x) { fprintf(stderr, "unknown type\n"); }
template<> void scan1(long& x) { if( scanf("%ld", &x) < 0 ) exit(0); }
void scan() {}
template<typename Head, typename... Tail>
void scan(Head& x, Tail&... xs) {
  scan1(x); scan(xs...);
}

template<typename W>
struct N008 {
   typedef std::vector<long> LI;
   typedef std::vector<W>    LW;
   long n, e;
   const LI &ss, &ds;
   const LW &ws;
   std::vector<LI> fi;
   std::vector<LI> ri;
   N008(long n_, const LI& ss_, const LI& ds_, const LW& ws_)
      : n(n_), e(ss_.size()), ss(ss_), ds(ds_), ws(ws_) {
      ri.resize(n+1);
      fi.resize(n+1);
      for(long i = 0; i < e; ++i) {
         fi[ss[i]].push_back(i);
         ri[ds[i]].push_back(i);
      }
   }
};

// N008 as undirected graph -> dfs tree
template<typename W>
struct N011 {
   typedef std::vector<long> LI;
   typedef std::vector<W>    LW;
   const N008<W>& g;
   LI ss, ds;
   LW ws;
   std::vector<bool> visited;
   N011(const N008<W>& g_, long t) : g(g_), visited(g.n + 1) {
      visited[t] = true;
      dfs(t);
   }
   void dfs(long i) {
      for(bool b : {true, false}) {
         for(long k : b ? g.fi[i] : g.ri[i]) {
            long to = b ? g.ds[k] : g.ss[k];
            if( visited[to] ) continue;
            visited[to] = true;
            ss.push_back(i);
            ds.push_back(to);
            ws.push_back(g.ws[k]);
            dfs(to);
         }
      }
   }
   N008<W> gen() {
      return N008<W>(g.n, ss, ds, ws);
   }
};

struct Solver {
   Solver() { fprintf(stderr, "--------Solver begin--------\n"); }
   ~Solver() { fprintf(stderr, "--------Solver end--------\n"); }
   void solve() {
      long v, x;
      scan(v, x);
      long e = v - 1;
      std::vector<long> xs(v+1, 0);
      for(long i = 1; i <= v; ++i) {
         scan(xs[i]);
      }
      std::vector<long> ss(e), ds(e), ws(e);
      for(long i = 0; i < e; ++i) {
         scan(ss[i], ds[i]);
      }
      N008<long> g(v, ss, ds, ws);
      N011<long> dfsMaker(g, x);
      N008<long> g2 = dfsMaker.gen();
      std::vector<bool> checked(v + 1, false);
      for(long i = 1; i <= v; ++i) {
         if( not ( xs[i] == 1 ) ) continue;
         long t = i;
         for(;;) {
            if( g2.ri[t].size() == 0 ) break;
            checked[t] = true;
            t = g2.ss[g2.ri[t][0]];
         }
      }
      long res = 0;
      for(long i = 1; i <= v; ++i) {
         if( checked[i] ) res += 2;
      }
      printf("%ld\n", res);
   }
};

int main() {
  loop std::unique_ptr<Solver>(new Solver())->solve();
}

Submission Info

Judge Result