fork download
  1. vertices, @max_cats_number = gets.split.map(&:to_i)
  2. @cats_positions = gets.split.map(&:to_i)
  3. @tree = Array.new(vertices + 1){ Array.new }
  4.  
  5. (vertices - 1).times do
  6. from, to = gets.split.map(&:to_i)
  7. from = from - 1
  8. to = to - 1
  9. @tree[from] << to
  10. @tree[to] << from
  11. end
  12.  
  13. def graph(current_vertice, last_vertice, cats_number)
  14.  
  15. sum = 0
  16.  
  17. if cats_number > @max_cats_number
  18. return 0
  19. elsif @tree[current_vertice].length == 1 && @tree[current_vertice][0] == last_vertice
  20. return 1
  21. end
  22.  
  23. for i in 0...@tree[current_vertice].length
  24. if @tree[current_vertice][i] != last_vertice
  25. if @cats_positions[@tree[current_vertice][i]] == 1
  26. sum = sum + graph(@tree[current_vertice][i], current_vertice, cats_number + 1)
  27. else
  28. sum = sum + graph(@tree[current_vertice][i], current_vertice, 0)
  29. end
  30. end
  31. end
  32.  
  33. return sum
  34. end
  35.  
  36. puts graph(0, -1, @cats_positions[0])
  37.  
Success #stdin #stdout 0s 28680KB
stdin
4 1
1 1 0 0
1 2
1 3
1 4
stdout
2