1class Candy:
  2
  3
  4    def __init__(self, mass, uranium):
  5        self.mass = mass
  6        self.uranium = uranium
  7
  8    def get_mass(self):
  9        return self.mass
 10
 11    def get_uranium_quantity(self):
 12        return self.uranium * self.mass
 13
 14def get_mass(candy):
 15    return candy.get_mass()
 16
 17def get_heaviest(candies):
 18    candies.sort(key = get_mass, reverse = True)
 19    return candies[0]
 20
 21class Person:
 22
 23
 24    def __init__(self, position):
 25        self.position = position
 26        self.x = position[0]
 27        self.y = position[1]
 28
 29    def get_position(self):
 30        return self.position
 31
 32    def get_x(self):
 33        return self.x
 34
 35    def get_y(self):
 36        return self.y
 37
 38    def set_position(self, position):
 39        self.position = position
 40        self.x = position[0]
 41        self.y = position[1]
 42    """
 43    We don't really care about precise measure on distance 
 44    only use it to impose order on the elements
 45    so (euclidean distance) ** 2 is still going to work fine
 46    """
 47
 48    def get_distance_squared(self, other):
 49        return ((self.x - other.x) ** 2) + ((self.y - other.y) ** 2)
 50
 51def get_x(person):
 52    return person.get_x()
 53
 54def get_y(person):
 55    return person.get_y()
 56
 57class Kid(Person):
 58
 59
 60    def __init__(self, position, initiative):
 61        super().__init__(position)
 62        self.initiative = initiative
 63        self.bag_of_candies = []
 64
 65    def get_initiative(self):
 66        return self.initiative
 67
 68    def add_candy(self, sweet):
 69        self.bag_of_candies.append(sweet)
 70
 71    def is_critical(self):
 72        uranium_totals = [
 73            self.bag_of_candies[i].get_uranium_quantity()
 74            for i in range(0, len(self.bag_of_candies))
 75        ]
 76        result = 0
 77        for total in uranium_totals:
 78            result = result + total
 79        return result - 20 > 0
 80
 81def get_initiative(kid):
 82    return kid.get_initiative()
 83
 84class Host(Person):
 85
 86
 87    def __init__(self, position, candies):
 88        super().__init__(position)
 89        self.bag_of_candies = [
 90            Candy(candies[i][0], candies[i][1])
 91            for i in range(0, len(candies))
 92        ]
 93
 94    def remove_candy(self, func):
 95        if not self.bag_of_candies:
 96            return None
 97        chosen_one = func(self.bag_of_candies)
 98        self.bag_of_candies.remove(chosen_one)
 99        return chosen_one
100
101"""
102A decorated function for the fixed distance relative to a specific kid
103"""
104def get_distance_squared(kid):
105    def get_fixed_distance(host):
106        return host.get_distance_squared(kid)
107    return get_fixed_distance
108
109class FluxCapacitor:
110
111
112    def __init__(self, participants):
113
114        self.hosts = []
115        self.kids = []
116        self.remaining_hosts = {}
117
118        for person in participants:
119            if isinstance(person, Kid):
120                self.kids.append(person)
121            else:
122                self.hosts.append(person)
123
124        # for each kid keep a personal record of hosts
125        for kid in self.kids:
126            copy_hosts = []
127            copy_hosts.extend(self.hosts)
128            self.remaining_hosts.update({kid : copy_hosts})
129
130        # sort the kids list according to the initiative descending
131        self.kids.sort(key = get_initiative, reverse = True)
132
133        # for each (kid, hosts) use the decorated distance function and sort
134        # according to the conditions in order of significance
135        for kid, hosts in self.remaining_hosts.items():
136            fixed_distance = get_distance_squared(kid)
137            hosts.sort(
138                key = lambda host:
139                (fixed_distance(host), get_x(host), get_y(host))
140            )
141
142    def get_victim(self):
143
144        dead =  set()
145
146        # worst case we go through all the hosts of each kid
147        for lap in range(0, len(self.hosts)):
148            for kid in self.kids:
149                # extract the candy from the first host of the kid in order
150                sweet = self.remaining_hosts[kid][0].remove_candy(get_heaviest)
151                if sweet is None:
152                    continue
153                kid.add_candy(sweet)
154                kid.set_position(self.remaining_hosts[kid][0].get_position())
155                self.remaining_hosts[kid].pop(0)
156                if kid.is_critical():
157                    dead.add(kid)
158            if dead:
159                return dead
160
161        return None