1from math import sqrt
  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
 22    def get_position(self):
 23        return self.position
 24    
 25    def set_position(self, new_position):
 26        self.position = new_position
 27
 28
 29class Kid(Person):
 30    
 31
 32    def __init__(self, position, initiative):
 33        super().__init__(position)
 34        self.initiative = initiative
 35        self.basket = []
 36        self.CRITICAL_LEVEL_URANIUM = 20
 37
 38    def get_initiative(self):
 39        return self.initiative
 40    
 41    def add_candy(self, candy):
 42        print(candy.mass, candy.uranium)
 43        self.basket.append(candy)
 44
 45    def is_critical(self):
 46        return sum(candy.get_uranium_quantity() for candy in self.basket) > self.CRITICAL_LEVEL_URANIUM
 47
 48
 49class Host(Person):
 50    def __init__(self, position, candies):
 51        super().__init__(position)
 52        self.candies = [Candy(mass, uranium) for mass,uranium in candies]
 53
 54    def remove_candy(self, func):
 55        if not self.candies:
 56            return None
 57
 58        selected_candy = func(self.candies)
 59        self.candies.remove(selected_candy)
 60        return selected_candy
 61
 62
 63class FluxCapacitor:
 64    def __init__(self, participants):
 65        self.kids = []
 66        self.hosts = []
 67
 68        for participant in participants: 
 69            if type(participant) == Kid:
 70                self.kids.append(participant)
 71            else:
 72                self.hosts.append(participant)
 73        
 74        self.kids = sorted(self.kids, key=lambda kid: kid.initiative, reverse=True)            
 75
 76    def get_victim(self):
 77        critical_kids = set()
 78        closest_host = None
 79        round = 0
 80        visited = [[False for _ in self.kids] for _ in self.hosts]
 81    
 82        while len(critical_kids) == 0 and round < len(self.hosts):
 83            for kid in self.kids:
 84                min_distance = float('inf')
 85                same_distance_hosts = []
 86                for host in self.hosts:
 87                    if visited[self.hosts.index(host)][self.kids.index(kid)]:
 88                        continue
 89                    distance = sqrt((kid.get_position()[0] - host.get_position()[0])**2 + (kid.get_position()[1] - host.get_position()[1])**2)
 90                    if distance < min_distance:
 91                        min_distance = distance
 92                        closest_host = host
 93                        same_distance_hosts = [host]
 94                    elif distance == min_distance:
 95                        same_distance_hosts.append(host)
 96                    if len(same_distance_hosts) > 1:
 97                        closest_host = min(same_distance_hosts, key=lambda host: host.get_position())
 98                
 99                visited[self.hosts.index(closest_host)][self.kids.index(kid)] = True
100                kid.set_position(closest_host.get_position())
101
102                if len(closest_host.candies) != 0: 
103                    kid.add_candy(closest_host.remove_candy(FluxCapacitor.find_max_element_in_list_by_criteria))
104
105                if kid.is_critical():
106                    critical_kids.add(kid)
107            round += 1
108
109        return critical_kids or None
110    
111
112    def find_max_element_in_list_by_criteria(list):
113        return max(list, key=lambda candy: candy.get_mass())