From 7dc67884d1ca499fc780ccad7aa05e9ebcd7d9e7 Mon Sep 17 00:00:00 2001 From: Caeden Perelli-Harris Date: Sun, 30 Jul 2023 19:38:09 +0100 Subject: [PATCH] Update cellular_automata/wa_tor.py Co-authored-by: Tianyi Zheng --- cellular_automata/wa_tor.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/cellular_automata/wa_tor.py b/cellular_automata/wa_tor.py index cb9252aed..7e7860fe4 100644 --- a/cellular_automata/wa_tor.py +++ b/cellular_automata/wa_tor.py @@ -208,14 +208,13 @@ class WaTor: prey_count, predator_count = len(prey), len(predators) - if prey_count > predator_count: - for entity in prey[:DELETE_UNBALANCED_ENTITIES]: - # Purge the first n entities of the prey - self.planet[entity.coords[0]][entity.coords[1]] = None - else: - for entity in predators[:DELETE_UNBALANCED_ENTITIES]: - # Purge the first n entities of the predators - self.planet[entity.coords[0]][entity.coords[1]] = None + entities_to_purge = ( + prey[:DELETE_UNBALANCED_ENTITIES] + if prey_count > predator_count + else predators[:DELETE_UNBALANCED_ENTITIES] + ) + for entity in entities_to_purge: + self.planet[entity.coords[0]][entity.coords[1]] = None def get_surrounding_prey(self, entity: Entity) -> list[Entity]: """