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, new_position):
 24        self.position = new_position
 25
 26
 27class Kid(Person):
 28    CRITICAL_MASS = 20
 29
 30    def __init__(self, position, initiative):
 31        super().__init__(position)
 32        self.initiative = initiative
 33        self.candies = []
 34
 35    def get_initiative(self):
 36        return self.initiative
 37
 38    def add_candy(self, candy):
 39        self.candies.append(candy)
 40
 41    def is_critical(self):
 42        return (
 43            sum(candy.get_uranium_quantity() for candy in self.candies)
 44            > Kid.CRITICAL_MASS
 45        )
 46
 47
 48class Host(Person):
 49    def __init__(self, position, candies):
 50        super().__init__(position)
 51        self.candies = [Candy(*t) for t in candies]
 52
 53    def remove_candy(self, function):
 54        if not self.candies:
 55            return None
 56
 57        candy_to_remove = function(self.candies)
 58        self.candies.remove(candy_to_remove)
 59
 60        return candy_to_remove
 61
 62
 63class FluxCapacitor:
 64    def __init__(self, participants):
 65        self.kids = [p for p in participants if isinstance(p, Kid)]
 66        self.hosts = [p for p in participants if isinstance(p, Host)]
 67        self.fallen = set()
 68
 69        self.unvisited_dict = {kid: set(self.hosts) for kid in self.kids}
 70
 71    @staticmethod
 72    def _get_max_mass_candy(candies):
 73        return max(candies, key=lambda c: c.mass)
 74
 75    def _get_distance(self, kid, host):
 76        return math.dist(kid.position, host.position)
 77
 78    def _get_closest_host(self, kid, hosts):
 79        return min(
 80            hosts,
 81            key=lambda host: (self._get_distance(kid, host), host.position),
 82        )
 83
 84    def _move_kid(self, kid):
 85        unvisited_hosts = self.unvisited_dict[kid]
 86
 87        if not unvisited_hosts:
 88            return None
 89
 90        closest_host = self._get_closest_host(kid, unvisited_hosts)
 91
 92        kid.set_position(closest_host.position)
 93        unvisited_hosts.remove(closest_host)
 94
 95        return closest_host
 96
 97    def _do_iteration(self):
 98        for kid in self.kids:
 99            host = self._move_kid(kid)
100            candy = host.remove_candy(FluxCapacitor._get_max_mass_candy)
101
102            if candy:
103                kid.add_candy(candy)
104
105            if kid.is_critical():
106                self.fallen.add(kid)
107
108    def get_victim(self):
109        iterations = len(self.hosts)
110        self.kids.sort(key=lambda kid: kid.initiative, reverse=True)
111
112        for _ in range(iterations):
113            self._do_iteration()
114
115            if self.fallen:
116                return self.fallen
117
118        return None