Validator for problem Jogurt

gistfile1.rb

#!/usr/bin/env ruby
def wrong_answer(reason)
  puts "Wrong Answer: #{reason}"
  exit(0)
end

# Checks if array is a permutation of the numbers 1, 2, 3, ..., 2 ** n - 1
def check_permutation(array, n)
  valid_permutation = (1..2 ** n - 1).to_a
  if array.size != valid_permutation.size
    wrong_answer "Team's tree has %d nodes, but expected %d." % [array.size, valid_permutation.size]
  end
  sorted_array = array.sort
  if sorted_array != valid_permutation
    missing_numbers = valid_permutation - sorted_array
    wrong_answer "The following numbers were not found on team's tree: %s" % [missing_numbers.join(", ")]
  end
end


# Checks if the difference between the sum of the two subtrees of tree equals 2 ** depth
# Also, recursively calls itself with the two subtrees of tree.
# Tree is assumed to be a complete binary tree sorted in preorder.
def check_sums(tree, depth)
  root = tree.first
  n = tree.size
  raise "Tree doesn't have an odd number of nodes" unless n % 2 == 1
  raise "Tree is not a complete binary tree" unless (n + 1) & n == 0 # Check that (n + 1) is a power of two. This is just an internal check.
  
  if n > 1
    left_subtree = tree[1..n/2]
    right_subtree = tree[n/2+1..n-1]
    
    left_sum = left_subtree.inject(0) { |sum, element| sum + element }
    right_sum = right_subtree.inject(0) { |sum, element| sum + element }
    
    expected_difference = 2 ** depth
    actual_difference = (left_sum - right_sum).abs
    if actual_difference != expected_difference
      wrong_answer "Difference for subtree rooted at #{root} should be #{expected_difference} but was #{actual_difference}. Left subtree's sum is #{left_sum} and right subtree's sum is #{right_sum}."
    end
    
    # Recursively check the rest of the tree
    check_sums(left_subtree, depth + 1)
    check_sums(right_subtree, depth + 1)
  end
end

#################################### MAIN ######################################

if ARGV.size < 3
  puts "Usage: #{__FILE__} <testdata.in> <program.out> <testdata.out>"
  exit 1
end
  
in_file = ARGV[0]
team_file = ARGV[1]
ans_file = ARGV[2]

n = File.read(in_file).to_i
tree = File.read(team_file).split.map(&:to_i)

check_permutation(tree, n)
check_sums(tree, 0)

# Accepted
exit(0)