1import math
  2
  3
  4class Candy:
  5
  6    def __init__(self, mass: int, uranium: float):
  7        self.mass = mass
  8        self.uranium = uranium
  9
 10    def get_uranium_quantity(self):
 11        return self.mass * self.uranium
 12
 13    def get_mass(self):
 14        return self.mass
 15
 16
 17class Person:
 18
 19    def __init__(self, position: tuple):
 20        self.position = position
 21
 22    def get_position(self):
 23        return self.position
 24
 25    def set_position(self, new_position: tuple):
 26        self.position = new_position
 27
 28
 29class Kid(Person):
 30
 31    def __init__(self, position: tuple, initiative: int):
 32        super().__init__(position)
 33        self.initiative = initiative
 34        self.candy_basket = []
 35        self.hosts_visited = set()
 36
 37    def get_initiative(self):
 38        return self.initiative
 39
 40    def add_candy(self, candy: Candy):
 41        self.candy_basket.append(candy)
 42
 43    def is_critical(self):
 44        current_uranium = 0
 45        candy: Candy
 46        for candy in self.candy_basket:
 47            current_uranium += candy.get_uranium_quantity()
 48
 49        if current_uranium > 20:
 50            return True
 51        return False
 52
 53    def get_hosts_visited(self):
 54        return self.hosts_visited
 55
 56    def visit_host(self, host):
 57        self.hosts_visited.add(host)
 58
 59
 60class Host(Person):
 61
 62    def __init__(self, position: tuple, candies: list):
 63        super().__init__(position)
 64        self.candies = []
 65        for candy_tuple in candies:
 66            self.candies.append(Candy(candy_tuple[0], candy_tuple[1]))
 67
 68    def remove_candy(self, filter_func):
 69        if not self.candies:
 70            return None
 71
 72        to_remove = filter_func(self.candies)
 73        if to_remove is None:
 74            return None
 75        self.candies.remove(to_remove)
 76
 77        return to_remove
 78
 79
 80class FluxCapacitor:
 81
 82    def __init__(self, participants: set):
 83
 84        self.kids = []
 85        self.hosts = []
 86        person: Person
 87        for person in participants:
 88            if type(person) == Kid:
 89                self.kids.append(person)
 90            else:
 91                self.hosts.append(person)
 92
 93        self.kids.sort(key=lambda x: x.initiative, reverse=True)
 94
 95    def get_victim(self):
 96        victims = set()
 97
 98        if len(self.kids) == 0:
 99            return None
100
101        loop = True
102
103        while loop:
104
105            kid: Kid
106            for kid in self.kids:
107                closest_host = self._find_closest_host_to_kid(kid)
108                if closest_host is None:
109                    loop = False
110                    break
111
112                given = closest_host.remove_candy(self.get_biggest_candy)
113                if given is not None:
114                    kid.add_candy(given)
115
116                if kid.is_critical():
117                    victims.add(kid)
118                    loop = False
119
120                kid.visit_host(closest_host)
121                kid.set_position(closest_host.get_position())
122
123
124        if victims:
125            return victims
126        return None
127
128    def _find_closest_host_to_kid(self, kid: Kid):
129        current_closest_dist = -1
130        closest_host = None
131
132        host: Host
133        for host in self.hosts:
134            curr_dist = math.dist(kid.get_position(), host.get_position())
135            if host in kid.get_hosts_visited():
136                continue
137
138            if current_closest_dist < 0 or curr_dist < current_closest_dist:
139                closest_host = host
140                current_closest_dist = curr_dist
141            elif curr_dist == current_closest_dist and host.get_position() < closest_host.get_position():
142                closest_host = host
143                current_closest_dist = curr_dist
144
145        return closest_host
146
147    def get_biggest_candy(self, host_basket: set):
148        biggest = None
149
150        candy: Candy
151        for candy in host_basket:
152            if biggest is None or biggest.get_mass() < candy.get_mass():
153                biggest = candy
154
155        return biggest
156
157    def get_kids(self):
158        return self.kids
159
160
161def get_biggest_candy(candies: list):
162    biggest = None
163
164    candy: Candy
165    for candy in candies:
166        if biggest is None or biggest.get_mass() < candy.get_mass():
167            biggest = candy
168
169    return biggest