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 the quantity of uranium in grams."""
 12        return self.mass * self.uranium
 13
 14    def get_mass(self):
 15        return self.mass
 16
 17
 18class Person:
 19
 20    def __init__(self, position):
 21        self.position = position
 22
 23    def get_position(self):
 24        return self.position
 25
 26    def set_position(self, position):
 27        self.position = position
 28
 29    @staticmethod
 30    def get_distance(person_1, person_2):
 31        """Calculate the distance between two people."""
 32        dx = person_1.position[0] - person_2.position[0]
 33        dy = person_1.position[1] - person_2.position[1]
 34        return math.sqrt(dx ** 2 + dy ** 2)
 35
 36
 37class Kid(Person):
 38
 39    CRITICAL_MASS = 20
 40
 41    def __init__(self, position, initiative):
 42        super().__init__(position)
 43        self.initiative = initiative
 44        self.basket = list()
 45
 46    def get_initiative(self):
 47        return self.initiative
 48
 49    def add_candy(self, candy_to_add):
 50        self.basket.append(candy_to_add)
 51
 52    def is_critical(self):
 53        """Sum the uranium contents of all candies and returns true if sum > 20."""
 54        uranium_sum = sum(candy.get_uranium_quantity() for candy in self.basket)
 55        return uranium_sum > Kid.CRITICAL_MASS
 56
 57
 58class Host(Person):
 59
 60    def __init__(self, position, candies):
 61        super().__init__(position)
 62        self.candies = [Candy(x, y) for x, y in candies]
 63        self.visited_from = []  # add children that already visited a given host here
 64        self.current_visitors = []  # add children that are visiting the host at the moment
 65
 66    def remove_candy(self, func_remove_candy):
 67        # if the host doesn't have candy, do nothing
 68        if not self.candies:
 69            return None
 70        # save the candy to delete from the list
 71        candy_to_remove = func_remove_candy(self.candies)
 72        self.candies.remove(candy_to_remove)
 73        return candy_to_remove
 74
 75
 76class FluxCapacitor:
 77
 78    def __init__(self, participants):
 79        self.kids = [person for person in participants if type(person) == Kid]
 80        self.hosts = [person for person in participants if type(person) == Host]
 81
 82    @staticmethod
 83    def find_candy_with_biggest_mass(candies):
 84        return max(candies, key=lambda some_candy: some_candy.mass)
 85
 86    def get_victim(self):
 87        while True:
 88            all_kids_visited_all_hosts = all(len(host.visited_from) == len(self.kids) for host in self.hosts)
 89            if all_kids_visited_all_hosts:
 90                break
 91
 92            # first, determine for each kid which house it's going to visit
 93            for kid in self.kids:
 94                try:
 95                    closest_unvisited_house = min([host for host in self.hosts if kid not in host.visited_from],
 96                                                  key=lambda host: (
 97                                                      Person.get_distance(kid, host), host.position[0],
 98                                                      host.position[1]))
 99                # ValueError will be thrown if the iterable is empty (there are no unvisited hosts for the current kid)
100                except ValueError:
101                    continue
102
103                # mark closest_unvisited_house as being actively visited by the respective kid
104                closest_unvisited_house.current_visitors.append(kid)
105                # the kid moves to the closest house, so change its coordinates
106                kid.set_position(closest_unvisited_house.position)
107
108            # secondly, make the hosts distribute the candy
109            for host in self.hosts:
110                while host.current_visitors:  # aim to empty the active visitors list
111                    # if the host doesn't have candy, remove all current visitors
112                    # and add them to the list of past visitors
113                    if not host.candies:
114                        host.visited_from += host.current_visitors
115                        host.current_visitors = []
116                        break
117
118                    # determine the next kid by comparing initiatives
119                    next_kid = max(host.current_visitors, key=lambda some_kid: some_kid.initiative)
120                    # find candy for the next kid
121                    next_candy = host.remove_candy(FluxCapacitor.find_candy_with_biggest_mass)
122                    # add candy to kid's basket
123                    next_kid.add_candy(next_candy)
124                    # save the kid in the list of past visitors and remove it from active visitors
125                    host.visited_from.append(next_kid)
126                    host.current_visitors.remove(next_kid)
127
128            # the last step is to add all children that have received critical mass to a set
129            irridated_kids = set()
130            for kid in self.kids:
131                if kid.is_critical():
132                    irridated_kids.add(kid)
133            if irridated_kids:
134                return irridated_kids
135
136        return None