1from collections import defaultdict
  2import math
  3
  4
  5class Candy:
  6    def __init__(self, mass, uranium):
  7        self.mass = mass
  8        self.uranium = uranium
  9
 10    def get_mass(self):
 11        return self.mass
 12
 13    def get_uranium_quantity(self):
 14        return self.mass * self.uranium
 15
 16
 17class Person:
 18    _basket = set()
 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, new_position):
 27        self.position = new_position
 28
 29
 30class Kid(Person):
 31    hosts_visited = set()
 32
 33    def __init__(self, position, initiative):
 34        super().__init__(position)
 35        self.initiative = initiative
 36
 37    def get_initiative(self):
 38        return self.initiative
 39
 40    def add_candy(self, candy):
 41        self._basket.add(candy)
 42
 43    def is_critical(self):
 44        return sum(candy.get_mass() for candy in self._basket) > 20
 45
 46
 47class Host(Person):
 48    def __init__(self, position, candies):
 49        super().__init__(position)
 50        self._basket = ({Candy(candy[0], candy[1]) for candy in candies})
 51
 52    def remove_candy(self, get_candy_to_give):
 53        if (len(self._basket)):
 54            candy_to_give = get_candy_to_give(self._basket)
 55            self._basket.remove(candy_to_give)
 56            return candy_to_give
 57        else:
 58            return None
 59
 60
 61class FluxCapacitor:
 62    def __init__(self, participants):
 63        self.victim = set()
 64        self.kids = set()
 65        self.hosts = set()
 66
 67        for participant in participants:
 68            if (isinstance(participant, Kid)):
 69                self.kids.add(participant)
 70            else:
 71                self.hosts.add(participant)
 72
 73        all_hosts_visited = False
 74
 75        while (not len(self.victim)) and (not all_hosts_visited):
 76            for kid in self.kids:
 77                hosts_to_visit = set()
 78                for host in self.hosts:
 79                    if (not host in kid.hosts_visited):
 80                        hosts_to_visit.add(host)
 81
 82                if (len(hosts_to_visit)):
 83                    nearest_host = (min(hosts_to_visit, key=lambda h: math.dist(
 84                        h.get_position(), kid.get_position())))
 85                    kid.set_position(nearest_host.get_position())
 86                    kid.hosts_visited.add(nearest_host)
 87                    candy_given = (
 88                        nearest_host.remove_candy(self._biggest_candy))
 89                    if (candy_given != None):
 90                        kid.add_candy(candy_given)
 91
 92            for kid in self.kids:
 93                if kid.is_critical():
 94                    self.victim.add(kid)
 95                if (not all_hosts_visited) and len(kid.hosts_visited) == len(self.hosts):
 96                    all_hosts_visited = True
 97
 98    def _kid_go_to_nearest_host(self, kid, hosts):
 99        hosts_to_visit = set()
100        for host in hosts:
101            if (host.get_position() != kid.get_position()) and (not host in kid.hosts_visited):
102                hosts_to_visit.add(host)
103
104        nearest_host = (min(hosts_to_visit, key=lambda h: math.dist(
105            h.get_position(), kid.get_position())))
106
107        kid.set_position(nearest_host.get_position())
108        kid.hosts_visited.add(nearest_host)
109
110        candy_given = nearest_host.remove_candy(self._biggest_candy)
111        if candy_given != None:
112            kid.add_candy(candy_given)
113
114    @staticmethod
115    def _biggest_candy(candies):
116        return max(candies, key=lambda c: c.get_mass())
117
118    def get_victim(self):
119        if len(self.victim):
120            return self.victim
121        return None