diff options
| -rw-r--r-- | robots/game.py | 42 |
1 files changed, 42 insertions, 0 deletions
diff --git a/robots/game.py b/robots/game.py index 827562a..cb065ec 100644 --- a/robots/game.py +++ b/robots/game.py @@ -168,6 +168,47 @@ class Game: for x, y, energy in robots ] + def resolve_combats(self): + to_remove = set() + to_append = defaultdict(list) + + for (x, y), robots in self.state.robots.items(): + if len(robots) <= 1: + # No combat here. + continue + + to_remove.add((x, y)) + + # each player's strength is the sum of their robots' energy + strength_by_player = defaultdict(int) + total = 0 + for player, energy in robots: + strength_by_player[player] += energy + total += energy + + winner_count = 0 + for player, strength in strength_by_player.items(): + # subtract the sum of the other players' strength + strength -= total - strength + # remove with strength players <= 0 + if strength > 0: + energy = min(strength, Energy.MAXIMUM) + to_append[player].append((x, y, energy)) + winner_count += 1 + + # we should end up with one player. + # (and I wrote a proof of that, so something's seriously wrong if + # this fails) + assert winner_count <= 1 + + for player, robots in self.state.robots_by_player.items(): + extra_robots = to_append[player] + robots[:] = [ + (x, y, energy) + for x, y, energy in robots + if (x, y) not in to_remove + ] + extra_robots + def __iter__(self): while not self.finished: yield self.state @@ -182,6 +223,7 @@ class Game: # 3. Factories produce new robots. self.factories_produce_robots() # 4. Combats are resolved. + self.resolve_combats() # 5. Each robot gains energy, up to the maximum. self.robots_gain_energy() |