Submission #1473719

---

Source Code Expand

// clang-format off
#include <bits/stdc++.h>
#define int long long
#define main signed main()
#define loop(i, a, n) for (int i = (a); i < (n); i++)
#define rep(i, n) loop(i, 0, n)
#define forever for (;;)
#define all(v) (v).begin(), (v).end()
#define rall(v) (v).rbegin(), (v).rend()
#define prec(n) fixed << setprecision(n)
template<typename A> using V = std::vector<A>;
template<typename A> using F = std::function<A>;
template<typename A, typename B> using P = std::pair<A, B>;
using pii = P<int, int>;
using vi = V<int>;
using vd = V<double>;
using vs = V<std::string>;
using vpii = V<pii>;
using vvi = V<vi>;
using vvpii = V<vpii>;
constexpr int INF = sizeof(int) == sizeof(long long) ? 1000000000000000000LL : 1000000000;
constexpr int MOD = 1000000007;
constexpr double PI = acos(-1);
template<typename A, typename B> bool cmin(A &a, const B &b) { return a > b ? (a = b, true) : false; }
template<typename A, typename B> bool cmax(A &a, const B &b) { return a < b ? (a = b, true) : false; }
template<typename T> std::istream &operator>>(std::istream &is, std::vector<T> &v) { for (T &x : v) is >> x; return is; }
template<typename A, typename B> std::istream &operator>>(std::istream &is, std::pair<A,B> &p) { is >> p.first; is >> p.second; return is; }
using namespace std;
// clang-format on

using Weight = int;
struct Edge {
  int src, dst;
  Weight weight;
  Edge(const int &s = 0, const int &d = 0, const Weight &w = 0) : src(s), dst(d), weight(w) {}
};

using Edges = std::vector<Edge>;
using Array = std::vector<Weight>;
using Matrix = std::vector<Array>;

class Graph {
  std::vector<Edges> g;
  using iterator = std::vector<Edges>::iterator;
  using const_iterator = std::vector<Edges>::const_iterator;

public:
  Graph(const int &size = 0) : g(size) {}
  int size() const { return g.size(); }
  const Edges &operator[](const int &i) const { return g[i]; }
  void addArc(const int &src, const int &dst, const Weight &w = 1) { g[src].emplace_back(src, dst, w); }
  void addEdge(const int &node1, const int &node2, const Weight &w = 1) {
    addArc(node1, node2, w);
    addArc(node2, node1, w);
  }
  iterator begin() { return g.begin(); }
  const_iterator begin() const { return g.begin(); }
  iterator end() { return g.end(); }
  const_iterator end() const { return g.end(); }
};

Graph toRootedTree(const Graph &g, const int &root = 0) {
  int n = g.size();
  Graph tree(n);
  std::vector<int> ord(n, -1);
  std::queue<int> q;
  q.push(root);
  int k = 0;
  ord[root] = k++;
  while (q.size()) {
    int u = q.front();
    q.pop();
    for (auto &e : g[u]) {
      int v = e.dst;
      if (ord[v] == -1) {
        ord[v] = k++;
        q.push(v);
        tree.addArc(u, v, e.weight);
      }
    }
  }
  return tree;
}

main {
  int n, x;
  cin >> n >> x;
  --x;
  vi h(n);
  cin >> h;
  Graph g(n);
  rep(_, n - 1) {
    int a, b;
    cin >> a >> b;
    g.addEdge(--a, --b);
  }
  g = toRootedTree(g, x);
  F<int(int)> dfs = [&](int v) {
    if (!g[v].size()) return 0LL;
    int c = 0;
    for (auto &e : g[v]) {
      int r = dfs(e.dst);
      c += r ? r + 2 : h[e.dst] * 2;
    }
    return c;
  };
  cout << dfs(x) << endl;
}

Submission Info

Judge Result