AtCoder ARC #030 B - ツリーグラフ

arc030-b.rb

n, x = gets.chomp.split(" ").map(&:to_i)
node = Struct.new("Node", :to)
@h = [0] + gets.chomp.split(" ").map(&:to_i)
@g = Array.new(n+1){node.new([])}
(n-1).times do
  a, b = gets.chomp.split(" ").map(&:to_i)
  @g[a].to << b
  @g[b].to << a
end

def jewel(v, root, from)
  cost = 0
  @g[v].to.each do |i|
    next if i == from
    cost += jewel(i, root, v)
  end
  # Root
  return cost if v == root
  # Leaf
  return cost + @h[v] if cost == 0
  # Other
  return cost + 1
end

puts jewel(x, x, -1) * 2