1import math
  2
  3class Candy:
  4    def __init__(self, mass: int, uranium: float):
  5        self.mass = mass
  6        self.uranium = uranium
  7
  8    def get_mass(self):
  9        return self.mass
 10
 11    def get_uranium_quantity(self):
 12        return self.mass * self.uranium
 13
 14
 15class Person:
 16    def __init__(self, position: tuple):
 17        self.position = position
 18
 19    def get_position(self):
 20        return self.position
 21
 22    def set_position(self, position: tuple):
 23        self.position = position
 24
 25    @staticmethod
 26    def dist_between(position1: tuple, position2: tuple):
 27        return math.sqrt((position2[0] - position1[0]) ** 2 + (position2[1] - position1[1]) ** 2)
 28
 29
 30class Kid(Person):
 31    CRITICAL_MASS_THRESHOLD = 20
 32
 33    def __init__(self, position: tuple, initiative: int):
 34        self.initiative = initiative
 35        self.basket = []
 36        super().__init__(position)
 37
 38    def get_initiative(self):
 39        return self.initiative
 40
 41    def add_candy(self, candy: Candy):
 42        self.basket.append(candy)
 43
 44    def is_critical(self):
 45        total_uranium = 0
 46        for candy in self.basket:
 47            total_uranium += candy.get_uranium_quantity()
 48        return total_uranium > self.CRITICAL_MASS_THRESHOLD
 49
 50
 51class Host(Person):
 52    def __init__(self, position: tuple, candies: list):
 53        self.basket = list()
 54        for candy in candies:
 55            self.basket.append(Candy(candy[0], candy[1]))
 56        super().__init__(position)
 57
 58    def remove_candy(self, predicate):
 59        if self.basket:
 60            candy = predicate(self.basket)
 61            self.basket.remove(predicate(self.basket))
 62            return candy
 63        return None
 64
 65
 66class FluxCapacitor:
 67    def __init__(self, participants: set):
 68        self.participants = participants
 69        self.hosts = [host for host in participants if isinstance(host, Host)]
 70        self.kids = [kid for kid in participants if isinstance(kid, Kid)]
 71
 72    def get_victim(self):
 73        hosts = self.hosts
 74        kids = self.kids
 75        kids_visited_hosts = {kid: [] for kid in kids}
 76
 77        while [len(visited_hosts) for visited_hosts in kids_visited_hosts.values()].count(len(hosts)) != len(kids):
 78            host_expected_kids = {host: [] for host in hosts}
 79
 80            # Determine which children will be visiting each host
 81            for kid in kids:
 82                closest_host = Host(position=(float('inf'), float('inf')), candies=[])
 83                for host in hosts:
 84                    if host not in kids_visited_hosts[kid]:
 85                        if (dist_a := Person.dist_between(kid.position, host.position)) < (
 86                                dist_b := Person.dist_between(kid.position, closest_host.position)):
 87                            closest_host = host
 88                        elif dist_a == dist_b:
 89                            if host.position[0] < closest_host.position[0]:
 90                                closest_host = host
 91                            elif host.position[0] == closest_host.position[0] and host.position[1] < closest_host.position[1]:
 92                                closest_host = host
 93                if closest_host.position != (float('inf'), float('inf')):
 94                    host_expected_kids[closest_host].append(kid)
 95
 96            # Sort kids by initiative
 97            for kids_list in host_expected_kids.values():
 98                if len(kids_list) > 1:
 99                    kids_list.sort(key=lambda x: x.initiative, reverse=True)
100
101            # Give candy to kids in defined order
102            for host in hosts:
103                # For every kid that wants to visit the host
104                for visiting_kid in host_expected_kids[host]:
105                    # Check if there is any candy left
106                    visiting_kid.set_position(host.get_position())
107                    if host.basket:
108                        # Host removes the candy from their basket and gives it to the kid
109                        visiting_kid.add_candy(host.remove_candy(lambda candy_basket: max(candy_basket, key=lambda candy: candy.get_mass())))
110                    # Kid now has visited this host and will never meet them again!
111                    kids_visited_hosts[visiting_kid].append(host)
112
113            # If there are any dead kids, return them
114            if dead_kids := set(kid for kid in kids if kid.is_critical()):
115                return dead_kids
116
117        return None