1import math
  2from collections import defaultdict
  3
  4
  5class Candy:
  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    def __init__(self, position):
 19        self._position = position
 20
 21    def set_position(self, position):
 22        self._position = position
 23
 24    def get_position(self):
 25        return self._position
 26
 27
 28class Kid(Person):
 29
 30    _BOUNDARY = 20
 31
 32    def __init__(self, position, initiative):
 33        super().__init__(position)
 34        self._initiative = initiative
 35        self._candies = []
 36
 37    def get_initiative(self):
 38        return self._initiative
 39
 40    def add_candy(self, new_candy):
 41        if new_candy:
 42            self._candies.append(new_candy)
 43
 44    def is_critical(self):
 45        return sum(map(lambda candy: candy.get_uranium_quantity(), self._candies)) > self._BOUNDARY
 46
 47
 48class Host(Person):
 49    def __init__(self, position, candies):
 50        super().__init__(position)
 51        self._basket = {Candy(*args) for args in candies}
 52
 53    def remove_candy(self, function):
 54        if not self._basket:
 55            return None
 56
 57        retrieved_candy = function(self._basket)
 58        self._basket.remove(retrieved_candy)
 59
 60        return retrieved_candy
 61
 62    def get_basket(self):
 63        return self._basket
 64
 65
 66class FluxCapacitor:
 67    def __init__(self, participants):
 68        self._hosts = {participant.get_position(): participant for participant in participants if isinstance(participant, Host)}
 69        self._kids = {participant: set() for participant in participants if isinstance(participant, Kid)}
 70
 71    def get_victim(self):
 72        dead_kids = set()
 73        roams = len(self._hosts)
 74
 75        for _ in range(roams):
 76            self._roam_around_pripyat()
 77
 78            for dead_kid in filter(lambda kid: kid.is_critical(), self._kids.keys()):
 79                dead_kids.add(dead_kid)
 80
 81            if dead_kids:
 82                return dead_kids
 83
 84        return None
 85
 86    def _roam_around_pripyat(self):
 87        hosts_and_their_kids = defaultdict(list)
 88
 89        for kid in sorted(self._kids.keys(), key=lambda current: current.get_initiative(), reverse=True):
 90            hosts_and_their_kids[self._get_closest_host(kid)].append(kid)
 91
 92        for host, their_kids in hosts_and_their_kids.items():
 93            for kid in their_kids:
 94                old_key = kid
 95
 96                kid.add_candy(host.remove_candy(lambda candies: max(candies, key=lambda candy: candy.get_mass())))
 97
 98                self._hosts[host.get_position()] = host
 99                self._kids[kid] = self._kids.pop(old_key)
100
101    def _get_closest_host(self, kid):
102        closest_distance = None
103        current_host = None
104
105        for host in self._hosts.values():
106            if host.get_position() in self._kids[kid]:
107                continue
108
109            distance = math.dist(host.get_position(), kid.get_position())
110            is_closest = False
111
112            if not current_host:
113                is_closest = True
114            else:
115                is_closest |= distance < closest_distance
116                is_closest |= (math.isclose(distance, closest_distance) and
117                               (host.get_position()[0] < current_host.get_position()[0]))
118                is_closest |= (math.isclose(distance, closest_distance) and
119                               host.get_position()[0] == current_host.get_position()[0] and
120                               host.get_position()[1] < current_host.get_position()[1])
121
122            if is_closest:
123                closest_distance = distance
124                current_host = host
125
126        self._kids[kid].add(current_host.get_position())
127        kid.set_position(current_host.get_position())
128
129        return current_host