1import math
  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
 17class Person:
 18
 19    def __init__(self, position):
 20        self.position = position
 21        self.basket = []
 22
 23    def get_position(self):
 24        return self.position
 25
 26    def set_position(self, new_position):
 27        self.position = new_position
 28
 29    def get_distance(self, person):
 30        return math.dist(self.position, person.get_position())
 31
 32
 33class Kid(Person):
 34
 35    CRITICAL_URANIUM_MASS = 20
 36
 37    def __init__(self, position, initiative):
 38        super().__init__(position)
 39        self.initiative = initiative
 40        self.visited_hosts = []
 41
 42    def get_initiative(self):
 43        return self.initiative
 44
 45    def add_candy(self, candy):
 46        self.basket.append(candy)
 47
 48    def is_critical(self):
 49        uranium_quantity = sum(candy.get_uranium_quantity() for candy in self.basket)
 50        return uranium_quantity > Kid.CRITICAL_URANIUM_MASS
 51
 52    def get_visited_hosts(self):
 53        return self.visited_hosts
 54
 55    def add_visited_host(self, host):
 56        self.visited_hosts.append(host)
 57
 58
 59class Host(Person):
 60
 61    def __init__(self, position, candies):
 62        super().__init__(position)
 63        self.basket = [Candy(mass, uranium) for mass, uranium in candies]
 64
 65    def remove_candy(self, candy_removal):
 66        if not self.basket:
 67            return None
 68        candy_to_remove = candy_removal(self.basket)
 69        self.basket.remove(candy_to_remove)
 70        return candy_to_remove
 71
 72
 73class FluxCapacitor:
 74
 75    def __init__(self, participants):
 76        self.participants = participants
 77        self.kids = [participant for participant in self.participants if isinstance(participant, Kid)]
 78        self.hosts = [participant for participant in self.participants if isinstance(participant, Host)]
 79
 80    def candy_to_remove(self, basket):
 81        return max(basket, key=lambda candy: candy.get_mass())
 82
 83    def get_victim(self):
 84        victims = ()
 85
 86        while self.kids:
 87            for kid in self.kids:
 88                if kid.is_critical():
 89                    victims.add(kid)
 90                    self.kids.remove(kid)
 91                else:
 92                    unvisited_hosts = [host for host in self.hosts if host not in kid.get_visited_hosts()]
 93                    if unvisited_hosts:
 94                        closest_host = min(unvisited_hosts, key=lambda host: kid.get_distance(host))
 95                        candy = closest_host.remove_candy(self.candy_to_remove)
 96                        if candy:
 97                            kid.add_candy(candy)
 98                            kid.add_visited_host(closest_host)
 99                            kid.set_position(closest_host.get_position())
100                    else:
101                        self.kids.remove(kid)
102        return victims or None