1import math
  2
  3
  4class Candy:
  5    def __init__(self, mass, uranium):
  6        self.mass = mass
  7        self.uranium = uranium
  8
  9    def get_uranium_quantity(self):
 10        return self.mass * self.uranium
 11
 12    def get_mass(self):
 13        return self.mass
 14
 15
 16class Person:
 17    def __init__(self, position):
 18        self.position = position
 19
 20    def get_position(self):
 21        return self.position
 22
 23    def set_position(self, position):
 24        self.postion = position
 25
 26
 27class Kid(Person):
 28    MAX_URANIUM_QUANTITY = 20
 29
 30    def __init__(self, position, initiative):
 31        super().__init__(position)
 32        self.initiative = initiative
 33        self.bag = []
 34
 35    def get_initiative(self):
 36        return self.initiative
 37
 38    def add_candy(self, candy):
 39        self.bag.append(candy)
 40
 41    def is_critical(self):
 42        return (
 43            sum([candy.get_uranium_quantity() for candy in self.bag])
 44            > Kid.MAX_URANIUM_QUANTITY
 45        )
 46
 47
 48class Host(Person):
 49    def __init__(self, position, candies):
 50        super().__init__(position)
 51        self.candies = [Candy(mass, uranium) for mass, uranium in candies]
 52
 53    def remove_candy(self, func):
 54        if self.candies:
 55            chosen_candy = func(self.candies)
 56            self.candies.remove(chosen_candy)
 57            return chosen_candy
 58
 59
 60class FluxCapacitor:
 61    def __init__(self, participants):
 62        self.kids = sorted(
 63            [participant for participant in participants if type(participant) is Kid],
 64            key=lambda kid: kid.get_initiative(),
 65            reverse=True,
 66        )
 67        # we are sorting the kids list so we dont have to worry about prioritizing the kids by initiative
 68        self.hosts = sorted(
 69            [participant for participant in participants if type(participant) is Host],
 70            key=lambda host: host.get_position(),
 71        )
 72        # we are sorting the hosts so we dont have to worry about the case where there are two host on equal distance from a kid
 73
 74    def _simulate_once(self, visited):
 75        dead_kids = set()
 76        for kid in self.kids:
 77            nearest_host = min(
 78                self.hosts,
 79                key=lambda host: (
 80                    host in visited[kid.get_initiative()],
 81                    math.dist(host.get_position(), kid.get_position()),
 82                ),
 83            )
 84            # because self.hosts is sorted we are sure that if there are more then one host at the same distance,
 85            # the first in order would have smaller x or y cordinates
 86
 87            kid.set_position(nearest_host.get_position())
 88            # dont know if its required to change the position
 89
 90            visited[kid.get_initiative()].add(nearest_host)
 91
 92            candy = nearest_host.remove_candy(
 93                lambda candies: max(candies, key=lambda candy: candy.get_mass())
 94            )
 95            if candy != None:
 96                kid.add_candy(candy)
 97
 98            if kid.is_critical():
 99                dead_kids.add(kid)
100        return dead_kids
101
102    def get_victim(self):
103        visited = {kid.get_initiative(): set() for kid in self.kids}
104        while any((len(values) < len(self.hosts) for values in visited.values())):
105            dead_kids = self._simulate_once(visited)
106            if dead_kids:
107                return dead_kids