1from math import dist
  2
  3
  4class Candy:
  5
  6    def __init__(self, mass, uranium):
  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    def __str__(self) -> str:
 17        return f"({self.mass}, {self.uranium})"
 18
 19
 20class Person:
 21
 22    def __init__(self, position):
 23        self.position = position
 24
 25    def get_position(self):
 26        return self.position
 27
 28    def set_position(self, position):
 29        self.position = position
 30
 31    def __hash__(self) -> int:
 32        return self.position.__hash__()
 33
 34
 35class Kid(Person):
 36
 37    I_am_glowing_in_the_dark_magic_number = 20
 38
 39    def __init__(self, position, initiative):
 40        Person.__init__(self, position)
 41        self.initiative = initiative
 42        self.candies = []
 43        self.visited_hosts = set()
 44
 45    def get_initiative(self):
 46        return self.initiative
 47
 48    def add_candy(self, candy):
 49        if candy is not None:
 50            self.candies.append(candy)
 51
 52    def is_critical(self):
 53        return sum([candy.get_uranium_quantity() for candy in self.candies]) > self.I_am_glowing_in_the_dark_magic_number
 54
 55    def get_closest_host(self, hosts):
 56        available_hosts = [
 57            host for host in hosts if host not in self.visited_hosts]
 58
 59        if not available_hosts:
 60            return None
 61
 62        closest_host = min(available_hosts, key=lambda host:
 63                           (dist(host.get_position(), self.get_position()),
 64                            *host.get_position())
 65                           )
 66
 67        self.visited_hosts.add(closest_host)
 68        self.set_position(closest_host.get_position())
 69        return closest_host
 70
 71    def __str__(self) -> str:
 72        candy_str = ""
 73        for candy in self.candies:
 74            candy_str += str(candy) + " "
 75
 76        hosts_str = ""
 77        for host in self.visited_hosts:
 78            hosts_str += str(host.get_position()) + " "
 79
 80        return f"Kid at {self.position} with initiative {self.initiative} and candies {candy_str}visited hosts {hosts_str}"
 81
 82
 83class Host(Person):
 84
 85    def __init__(self, position, candies):
 86        Person.__init__(self, position)
 87        self.candies = [Candy(*candy) for candy in candies]
 88
 89    def remove_candy(self, func):
 90        if len(self.candies) == 0:
 91            return None
 92
 93        candy = func(self.candies)
 94        self.candies.remove(candy)
 95        return candy
 96
 97    def give_away_happiness(self, kids, func):
 98        for kid in sorted(kids, key=lambda kid: kid.get_initiative(), reverse=True):
 99            kid.add_candy(self.remove_candy(func))
100
101
102class FluxCapacitor:
103
104    def __init__(self, participants):
105        self.kids = [participant for participant in participants if isinstance(
106            participant, Kid)]
107        self.hosts = [participant for participant in participants if isinstance(
108            participant, Host)]
109
110    def get_victim(self):
111
112        # RIP dead_kids_tell_no_tales() 05.11.2023 - 07.11.2023 :pepesad:
113        while True:
114            kids_to_hosts = {
115                host: [] for host in self.hosts
116            }
117
118            for kid in self.kids:
119                host = kid.get_closest_host(self.hosts)
120                if host is not None:
121                    kids_to_hosts[host].append(kid)
122
123            self.are_there_candies = False
124
125            def compare(candies):
126                self.are_there_candies = True
127                return max(candies, key=lambda candy: candy.get_mass())
128
129            for host in self.hosts:
130                host.give_away_happiness(kids_to_hosts[host], compare)
131
132            if not self.are_there_candies:
133                return None
134
135            dead_kids = {kid for kid in self.kids if kid.is_critical()}
136
137            if dead_kids:
138                return dead_kids