1import math
  2
  3class Candy:
  4
  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
 18    def __init__(self, position):
 19        self.position = position
 20
 21    def get_position(self):
 22        return self.position
 23    
 24    def set_position(self, new_position):
 25        self.position = new_position
 26
 27
 28class Kid(Person):
 29
 30    CRITICAL_URANIUM_QUANTITY_PER_KID = 20
 31
 32    def __init__(self, position, initiative):
 33        super().__init__(position)
 34        self.initiative = initiative
 35        self.uranium_quantity_in_basket = 0
 36
 37    def get_initiative(self):
 38        return self.initiative
 39    
 40    def add_candy(self, candy):
 41        self.uranium_quantity_in_basket += candy.get_uranium_quantity()
 42
 43    def is_critical(self):
 44        return self.uranium_quantity_in_basket > Kid.CRITICAL_URANIUM_QUANTITY_PER_KID
 45    
 46
 47class Host(Person):
 48
 49    def __init__(self, position, candies):
 50        super().__init__(position)
 51        real_candies = []
 52        for candy in candies:
 53            real_candies.append(Candy(candy[0], candy[1]))
 54        self.candies = sorted(real_candies, key=lambda candy: candy.mass, reverse=True)
 55        self.guests = []
 56
 57    def remove_candy(self, candy_to_be_removed):
 58        if self.candies:
 59            exact_candy_to_be_removed = candy_to_be_removed(self.candies)
 60            return exact_candy_to_be_removed
 61        
 62        return None
 63    
 64    def add_guest(self, guest):
 65        self.guests.append(guest)
 66        
 67    def treat(self):
 68        self.guests.sort(key=lambda kid: kid.initiative)
 69        kids_with_critical_condition = set()
 70
 71        for (kid, candy) in zip(self.guests, self.candies):
 72            kid.add_candy(candy)
 73            self.candies.remove(candy)
 74            kid.set_position(self.get_position())
 75
 76            if kid.is_critical():
 77                kids_with_critical_condition.add(kid)
 78
 79        self.guests.clear()
 80        return kids_with_critical_condition
 81
 82class FluxCapacitor:
 83
 84    def __init__(self, participants):
 85        self.hosts = []
 86        self.kids = []
 87
 88        for participant in participants:
 89            if type(participant) == Host:
 90                self.hosts.append(participant)
 91            else:
 92                self.kids.append(participant)
 93
 94    def arrange_guests(self):
 95        for kid in self.kids:
 96            lowest_distance = math.dist(kid.get_position(), self.hosts[0].get_position())
 97            closest_host = self.hosts[0]
 98
 99            for host in self.hosts:
100                current_distance = math.dist(kid.get_position(), host.get_position())
101
102                if current_distance < lowest_distance:
103                    lowest_distance = current_distance
104                    closest_host = host
105                elif current_distance == lowest_distance:
106                    if host.get_position()[0] < closest_host.get_position()[0] or host.get_position()[1] < closest_host.get_position()[0]:
107                        lowest_distance = current_distance
108                        closest_host = host
109
110            closest_host.add_guest(kid)
111
112    def get_victim(self):
113        self.arrange_guests()
114        final_set_critical_condition = set()
115        
116        for host in self.hosts:
117            final_set_critical_condition.update(host.treat())
118
119        return final_set_critical_condition