Facebook Careers Puzzle—swarm solution in Ruby

From Facebook’s Job Puzzles page:

You are the Zerg Queen in command of the 654195331th legion of Zerg forces, the only side worth playing in StarCraft. The wise Overmind has given you the important task of cleaning up some remote planets for mineral mining operations. Apparently, some pesky Terran forces have decided to set up base defenses in these locations, prior to your arrival. With your limited forces, you must determine which Terran bases to attack. Your Zerg forces on each planet have free movement over that one planet, and may split up to attack any number of Terran bases on that planet. However, your Zerg ground forces have not yet evolved space flight and thus cannot travel from planet to planet to assist each other.

Odds of victory over a particular Terran base are equal to: P(z,s) = e(-63s+10)/(e(-63s+10)+e(-21z))… etc

This problem is hard. On the surface, it seems like an optimization problem of two variables: minerals vs zerg, and then on top of that, deciding on which bases to attack. The equation above may look daunting, but with some pretty basic pre-calc, can be simplified. Look at the comments in my solution for the steps I took.

The first problem is deciding whether a base is worth attacking, and how many zerg are needed. That’s not the end of the story, as you can actually attack each base with different amounts of zerg to get different amounts of minerals. It’s a cluster..bleep. Thank goodness for the exp(), as it drastically limits the possible amount of forces to use for each base. I had to use a trick to make it work for large minerals. Can you spot it and figure out why it works?

After that, it’s just another 0-1 knapsack problem to solve, with thankfully only a very few amount of possibilities to consider.

For testing, I used forum David’s test generator to make this file (4MB).

    time ./swarm swarm.txt > out

    real    0m21.126s
    user    0m20.167s
    sys 0m0.107s

    MD5 (out) = 14fc14d34d05d485e1e261a86dd5a499

The following solution in Ruby passed the Facebook puzzle robot:

(swarm.rb) download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

#!/usr/bin/ruby
# ZERGRUSH
# Kanwei Li, 2009
#
# e^(-63s+10)/(e^(-63s+10)+e^(-21z)) * m >= 1
# m * e^(-63s+10) >= e^(-63s+10) + e^(-21z)
# m * e^(-63s+10) - e^(-63s+10) >= e^(-21z)
# e^(-63s+10) (m-1) >= e^(-21z)
# ln(m-1) + (-63s+10) >= -21z

# e^(-63s+10)/(e^(-63s+10)+e^(-21z)) * m = g
# m * e^(-63s+10) = g(e^(-63s+10) + e^(-21z))
# m * e^(-63s+10) - g*e^(-63s+10) = g*e^(-21z)
# e^(-63s+10) (m-g) = g*e^(-21z)
# ln(m-g) + (-63s+10) = ln(g) - 21z
# ln(g) - ln(m-g) = -63s + 10 + 21z
# g/(m-g) = e^(-63s + 10 + 21z)
# let cc = e^(-63s + 10 + 21z)
# g = (m-g) * cc
# g + g*cc = m*cc
# g(1 + cc) = m*cc
# g = m*cc/(1 + cc)

def optimize(ary, total)
  return [] if ary.empty?
  table = []
  (ary.size+1).times { |i| table[i] = [] }
  (0..total).each { |zerg| table[0][zerg] = 0 }
  (1..ary.size).each do |base|
    table[base][0] = 0
    (1..total).each do |zerg|
      if ary[base-1].zerg <= zerg && (ary[base-1].minerals + table[base-1][zerg - ary[base-1].zerg] > table[base-1][zerg])
        table[base][zerg] = ary[base-1].minerals + table[base-1][zerg - ary[base-1][1]]
      else
        table[base][zerg] = table[base-1][zerg]
      end
    end
  end
  result = []
  i, k = ary.size, total
  while i > 0 && k > 0
    if table[i][k] != table[i-1][k]
      result << ary[i-1]
      k -= ary[i-1].zerg
    end
    i -= 1
  end
  result
end

def cartesian_prod(ary)
  ary.inject([[]]){|old,lst|
    new = []
    lst.each{|e| new += old.map{|c| c.dup << e }}
    new
  }
end

def gain(minerals, terran, zerg)
  n = -63 * terran + 10 + 21 * zerg
  return 0 if n < 0
  return minerals if n > 10
  cc = Math.exp(n)
  (minerals * cc / (1 + cc)).round
end

def zerg_needed_for_feasible(minerals, terran)
  ((Math.log(minerals-1) + (-63 * terran + 10)) / -21).ceil
end

def feasible?(minerals, terran, zerg)
  return false if minerals <= 1
  Math.log(minerals-1) + (-63 * terran + 10) >= (-21 * zerg).to_f
end

# Make sure input file exists and read from it
filename = ARGV[0]
unless filename && File.exist?(filename)
  puts "error: must specify a valid input file"
  exit
end
input = File.open(filename)

num_planets = input.gets.to_i
out = ""

Base = Struct.new(:index, :zerg, :minerals)

num_planets.times do
  bases, zerg = input.gets.split(/\s+/).collect { |i| i.to_i }
  feasible = []
  bases.times do |n|
    terran, minerals = input.gets.split(/\s+/).collect { |i| i.to_i }
    if feasible?(minerals, terran, zerg)
      zerg_needed = zerg_needed_for_feasible(minerals, terran)
      last_gain = nil
      base_ary = []
      loop do
        g = gain(minerals, terran, zerg_needed)
        break if last_gain && last_gain == g
        base_ary << Base.new(n, zerg_needed, g)
        last_gain = g
        zerg_needed += 1
      end
      feasible << base_ary
    end
  end
  candidates = cartesian_prod(feasible)

  best_zerg, best_minerals, best_ary = nil, nil, nil
  candidates.each do |cand|
    cand = cand.sort_by { |a| [a.zerg, a.minerals] }
    cand_ary = optimize(cand, zerg)
    cand_ary.sort! { |a, b| a.index <=> b.index }
    total_zerg = cand_ary.inject(0) { |sum, i| sum += i.zerg }
    total_minerals = cand_ary.inject(0) { |sum, i| sum += i.minerals }
    if best_minerals.nil? || (total_minerals >= best_minerals)
      best_zerg, best_minerals = total_zerg, total_minerals
      best_ary = cand_ary
    end
  end
  out << "#{best_zerg} #{best_minerals}\n"
  out << best_ary.inject([]) { |ary, base| ary << base.index << base.zerg }.join(' ') << "\n"
end

puts out
Posted on