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.uranium*self.mass
 12    
 13    def get_mass(self):
 14        return self.mass
 15    
 16    def __gt__(self, other):
 17        return self.mass > other.mass
 18    
 19
 20class Person:
 21
 22    def __init__(self, position):
 23        self.set_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
 32class Kid(Person):
 33
 34    CRITICAL_URANIUM_DOSE = 20
 35
 36    def __init__(self, position, initiative):
 37        super().__init__(position)
 38        self.initiative = initiative
 39        self.basket = []
 40        self.visited_hosts = set()
 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        sum = 0
 50        for candy in self.basket:
 51            sum += candy.get_uranium_quantity()
 52        return sum > self.CRITICAL_URANIUM_DOSE
 53
 54    def __lt__(self, other):
 55        return self.initiative < other.initiative        
 56
 57    def _sort_hosts(self, host):
 58        dist = math.sqrt((self.position[0] - host.position[0]) ** 2 + (self.position[1] - host.position[1]) ** 2)
 59        return dist, host.position[0], host.position[1]
 60    
 61    def _go_to_door(self, hosts):
 62        unvisited_hosts = [host for host in hosts if host not in self.visited_hosts]
 63        if not unvisited_hosts:
 64            return False
 65        unvisited_hosts.sort(key=lambda host: self._sort_hosts(host))
 66        selected_host = unvisited_hosts[0]
 67        self.visited_hosts.add(selected_host)
 68        self.position=selected_host.position
 69        selected_host.visiting_kids.append(self)
 70
 71
 72class Host(Person):
 73
 74    def __init__(self, position, candies):
 75        super().__init__(position)
 76        self.basket = []
 77        self.visiting_kids = []
 78        for candy in candies:
 79            self.basket.append(Candy(candy[0], candy[1]))
 80
 81    def remove_candy(self, func):
 82        if not self.basket:
 83            return 
 84        candy_to_remove = func(self.basket)
 85        self.basket = [candy for candy in self.basket if candy is not candy_to_remove]
 86        return candy_to_remove
 87    
 88    def _give_candy(self):
 89        self.visiting_kids.sort(reverse = True)
 90        for kid in self.visiting_kids:
 91            candy = self.remove_candy(func=lambda basket:sorted(basket)[0])
 92            if candy is not None:
 93                kid.basket.append(candy)
 94        self.visiting_kids.clear()
 95
 96
 97class FluxCapacitor:
 98
 99    def __init__(self, participants):
100        self.participants = participants
101        self.kids = []
102        self.hosts = []
103        for person in participants:
104            if isinstance(person, Kid):
105                self.kids.append(person)
106            elif isinstance(person, Host):
107                self.hosts.append(person)
108    
109    def get_victim(self):
110        kids_with_critical_mass = set()
111        while True:
112            for kid in self.kids:
113                kid._go_to_door(self.hosts)
114            for host in self.hosts:
115                host._give_candy()
116            for kid in self.kids:
117                if kid.is_critical():
118                    kids_with_critical_mass.add(kid)
119            if kids_with_critical_mass:
120                return kids_with_critical_mass
121            all_kids_visited_all_hosts = True
122            for kid in self.kids:
123                if not kid.visited_hosts == set(self.hosts):
124                    all_kids_visited_all_hosts = False
125                    break
126            if all_kids_visited_all_hosts:
127                return None