Ruby: "MasterMind" solver

mm.rb

#!/usr/bin/env ruby

require 'set'

class GameBoard

    attr_reader :boxes, :colors, :scheme

    def initialize(colors, boxes, scheme = nil)
        raise ArgumentError.new("scheme has #{scheme.length} elements; expecting #{boxes}") if scheme && scheme.length != boxes
        @boxes = boxes
        @colors = colors
        @scheme = scheme || boxes.times.collect { Random.rand(colors) }
    end

    # Return an array of values 0, 1, or 2.  For each position:
    # 0: Wrong color
    # 1: Right color, wrong position
    # 2: Right color, right position
    # Returns nil if try_scheme matches the actual scheme.
    def test(try_scheme)
        raise ArgumentError.new("scheme has #{scheme.length} elements; expecting #{boxes}") unless try_scheme.length == @boxes
        return nil if try_scheme == @scheme
        ret = [nil] * @boxes
        try_scheme.each_with_index { |e, i| ret[i] = 2 if e == @scheme[i] }
        try_scheme.each_with_index do |e, i|
            next if ret[i]
            @scheme.each_with_index do |se, si|
                next if ret[si] == 2
                if se == e
                    ret[i] = 1
                    break
                end
            end
            ret[i] = 0 unless ret[i]
        end
        # p [@scheme, try_scheme, ret]
        ret
    end

end

def print_stats(s_tries, s_moves)
    count = s_tries.reduce(:+)
    colors = s_tries.length + 1
    t_moves = 0
    s_moves.each_with_index { |x, i| t_moves += x * i }
    col1_width = [count, colors].max.to_s.length
    col2_width = [count, t_moves].max.to_s.length

    puts
    puts "Tries:"
    sum = 0
    s_tries.each_with_index do |s, i|
        next if i.zero?
        printf "    %#{col1_width}d: %#{col2_width}d (%5.2f%%)\n", i, s, (100.0 * s / count)
        sum += s * i
    end
    printf "Average tries over %d attempts: %.2f\n\n", count, (0.0 + sum) / count

    puts
    puts "Moves:"
    sum = 0
    s_moves.each_with_index do |s, i|
        next if s.zero?
        printf "    %#{col1_width}d: %#{col2_width}d (%5.2f%%)\n", i, s, (100.0 * s * i / t_moves)
        sum += s * i
    end
    printf "Average moves over %d attempts: %.2f\n\n", count, (0.0 + sum) / count
end

#
# Solving algorithm:
#
#  1. Start with a guess that has as many colors as possible.
#  2. Consider all colors available to each individual position.
#  3. Exit if solved.
#  4. "Lock in" any correct items.
#  5. Remove any "wrong" items from the possible colors for all positions.
#  6. Remove any "right, but wrong place" colors from list of available
#     colors for their respective positions.
#  7. Find a position for each color in the set of "right, but wrong place".
#  8. Create a list of available colors, consisting of the aggregation of
#     available colors for each position.
#  9. In that list, deprioritize any colors that were used in step 7.
# 10. Populate the remaining positions using the first available color in
#     each position, from the list in step 9.
# 11. Go to step 3.
#
def solve(game_board)
    boxes = game_board.boxes
    colors = game_board.colors
    # available: sets of colors that are available to each position.
    available = boxes.times.collect { Set.new(0...colors) }
    # try_scheme: the solution to try.  Start with a solution that
    # represents as many colors as possible.
    try_scheme = boxes.times.collect { |x| x % colors }
    tries = 1
    moves = boxes
    while (res = game_board.test(try_scheme))
        tries += 1
        # must_use: colors that appear in the solution, but elsewhere.
        must_use = Set.new
        res.each_with_index do |e, i|
            c = try_scheme[i]
            case e
            when 0
                available.select { |a| a.instance_of?(Set) }.each { |a| a.delete(c) }
                moves += 1
            when 1
                available[i].delete(c)
                must_use << c
                moves += 1
            when 2
                available[i] = c
            end
        end
        # Start by placing the ones we know.
        try_scheme = available.map { |e| e.instance_of?(Set) ? nil : e }
        # Next, use everything in must_use
        while true
            mu = must_use.clone
            ovl = []
            boxes.times.reject { |i| try_scheme[i] }.sort { |a, b| available[a].length <=> available[b].length }.each do |i|
                c = (mu & available[i]).to_a.sample
                mu.delete(c)
                ovl[i] = c
            end
            # If we've exhausted must_use, then overlay the ones we've picked
            # onto try_scheme.
            if mu.length.zero?
                ovl.each_with_index { |e, i| try_scheme[i] ||= e }
                break
            end
        end
        # may_use is the aggregation of all available colors.
        may_use = available.select { |x| x.instance_of? Set }.reduce(:+)
        # For all positions that don't yet have values, use what's in
        # may_use, but prioritize anything that wasn't in must_use
        # (because they already appear elsewhere).
        use_order = (may_use - must_use).to_a + must_use.to_a
        boxes.times.reject { |i| try_scheme[i] }.each do |i|
            c = use_order.find { |x| available[i].include? x }
            # Use that color, and then push it to the back of the list.
            try_scheme[i] = c
            use_order.delete c
            use_order.push c
        end
    end
    return [tries, moves]
end

unless ARGV.length.between?(2,3)
    $stderr.puts "Usage: #{__FILE__} <colors> <boxes> [count]"
    exit 1
end

COLORS, BOXES, COUNT = ARGV.map { |i| i.to_i }.push(500)

s_tries = [0] * (COLORS + 1)
s_moves = [0] * (COLORS * BOXES + 1)
t_moves = 0

COUNT.times do
    g = GameBoard.new(COLORS, BOXES)
    tries, moves = solve(g)
    s_moves[moves] += 1
    s_tries[tries] += 1
    t_moves += moves
end

print_stats s_tries, s_moves
exit 0